Solved by ptrblck in post #2 
              
              
                In your current forward implementation you are passing x to F.softmax so that the actual model will not be used. 
Also, nn.CrossEntropyLoss expects raw logits, as internally F.log_softmax will be applied, so remove the F.softmax and return out directly. 
If you are still facing a shape mismatch erro…