Hi!
In your self.classification layer you use:
nn.Softmax(dim=1)

However the CrossEntropyLoss  from pytorch already does this, as the documentation states:

The input is expected to contain raw, unnormalized scores for each class

Also take a look at the mathematical formula, you’ll see that the softmax activation is already in there.
Does removing the softmax activation solve your problem?