Solved by ptrblck in post #5 
              
              
                You are currently calculating the max in the batch dimension, which seems to be wrong as you would usually compute the predictions by taking the argmax in the class dimension. 
This should work: 
output = model(data).squeeze() # output.shape == [12, 4]
preds = torch.argmax(output, dim=1) # preds.sha…