Solved by ptrblck in post #4 
              
              
                Your target tensor might have an additional unnecessary dimension as [batch_size, 1], so remove dim1 via target = target.squeeze(1) if that’s the case.