Solved by KFrank in post #2 
              
              
                Hello Carmelo! 

You can do anything you want in your test phase, regardless of how 
you trained your model.  (Just make sure that your test-phase results 
mean what you want them to mean.) 
You say you use CrossEntropyLoss.  The output of your model 
should therefore be raw-score logits, and these …