I think you should also pass the predictions as a hot encoded tensor.
If you pass the argmax as the prediction, it will be internally broadcasted to something like this (line of code 44):
num_features, num_hidden, num_classes = 5, 20, 10
batch_size = 10

features = torch.randn(batch_size, num_features, requires_grad=True)
targets = torch.randint(0, num_classes, (batch_size, ))
targets = F.one_hot(targets, num_classes=num_classes)

model = MultilayerPerceptron(num_features, num_hidden, num_classes)

logits, probas = model(features)

# Internally what F.mse_loss will do
ret = (torch.argmax(probas, 1).float() - targets.float()) ** 2

print(targets)
> tensor([[1, 0, 0, 0, 0, 0, 0, 0, 0, 0]])
print(probas)
> tensor([[0.4905, 0.4872, 0.5295, 0.5871, 0.3829, 0.4792, 0.5026, 0.5062, 0.4536,
         0.5444]], grad_fn=<SigmoidBackward>)
print(torch.argmax(probas, 1).float())
> tensor([3.])
print(ret)
> tensor([[4., 9., 9., 9., 9., 9., 9., 9., 9., 9.]])  # !!!

# Without the argmax it should be working I think
ret = (probas - targets.float()) ** 2
print(ret)
> tensor([[0.2596, 0.2374, 0.2803, 0.3447, 0.1466, 0.2296, 0.2526, 0.2563, 0.2058,
         0.2964]], grad_fn=<PowBackward0>)

Also, I’m not sure if you update the parameters manually, but if you would like to use an optimizer, I had to wrap all internal parameters in nn.Parameter.
Let me know, if this helps or if I’m misunderstanding something.