Solved by ptrblck in post #2 
              
              
                nn.CrossEntropyLoss expects raw logits as the model outputs, so you should remove the last softmax layer in your model. 
Do you have any other doubts regarding a specific part of the code?