Not working model

pytorch.py

import numpy as np
import torch
import torch.nn as nn
import torchvision.transforms as transforms
from torchvision import datasets

torch.manual_seed(1)


class MnistModel(nn.Module):
    def __init__(self, batch_size):
        super(MnistModel, self).__init__()
        self.batch_size = batch_size
        self.w = torch.nn.Parameter(torch.empty(batch_size, 784, 10).uniform_(0, 1))
        self.b = torch.nn.Parameter(torch.empty(10).uniform_(0, 1))

    def forward(self, x):
        return torch.bmm(x, self.w) + self.b


batch_size = 1000

classes = 10
train_data = datasets.MNIST('data', train=True,
                            download=True,
                            transform=transforms.ToTensor())
train_loader = torch.utils.data.DataLoader(train_data,
                                           batch_size=batch_size,
                                           shuffle=False)
test_data = datasets.MNIST('data', train=False, transform=transforms.ToTensor())
test_loader = torch.utils.data.DataLoader(test_data, batch_size=batch_size)

learning_rate = 0.5


model = MnistModel(batch_size)
print("weight[0][0][0] before training:", list(model.parameters())[0][0][0][0])
print("Bias[0] before training:", list(model.parameters())[1][0])

criterion = nn.MSELoss()
optimizer = torch.optim.SGD(model.parameters(), lr=learning_rate)
rows = np.array(range(batch_size))
zeros = torch.zeros((batch_size, 1, 10), dtype=torch.float64)

for i, (raw_data, raw_target) in enumerate(train_loader):
    data = raw_data.view((batch_size, 1, 784))

    logits = model(data)
    zeros[:, :, :] = 0
    zeros[rows, 0, raw_target] = 1
    loss = criterion(logits, zeros.float())

    # Backward and optimize
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()


print("\nweight[0][0][0] after training: ", list(model.parameters())[0][0][0][0])
print("Bias[0] after training", list(model.parameters())[1][0])
#RESULT:
#> weight[0][0][0] before training: tensor(0.7576)
#> Bias[0] before training: tensor(0.3681)
#>
#> weight[0][0][0] after training:  tensor(0.7576)
#> Bias[0] after training tensor(-48.1940)