It looks like you’re using F.nll_loss() but you’re returning raw unnormalized scores from your model. You either need to:
(a) use a cross entropy loss (either the one you’re assigning to criterion and never using again, or F.cross_entropy()); or
(b) continue to use the NLL loss but use nn.LogSoftmax() (or F.log_softmax()) as the last layer of your model. The NLL loss expects log probabilities as its input so you have to transform your unnormalized scores with a (log) softmax.