Solved by KFrank in post #2 
              
              
                Hi jpj! 

Yes, you should use CrossEntropyLoss as your loss criterions (and not 
use Softmax because CrossEntropyLoss expects logits rather than 
probabilities as its predicted input). 
Best. 
K. Frank