For the classification problem, e.g. MNIST, MSE loss should be avoided. Instead, the suitable loss should be CrossEntropyLoss (or a combination of nll_loss with soft_max). That’s why the first case yields not very bad accuracy.



 1337mathster:

x = self.fc(x).relu()


Also, remove relu here. It will squash the predicted logits into 0 if negative.