nn.CrossEntropyLoss expects logits while your model seems to apply a softmax on the output.
While this shouldn’t be responsible for the NaNs, you still should remove it as this might slow down your training.
To debug further, you could use Anomaly Detection 202 and try to see where these NaN values got produced.
Based on your description it seems like you are seeing NaN values in the output after one epoch. Is that correct?
Just for the sake of debugging, could you set drop_last=True in your DataLoader and run the code again?