Solved by thadar in post #4 
              
              
                if you want the probability for each class then yes. 
if you want only the predicted class then you don’t need it you can only use torch.max()