Solved by ptrblck in post #5 
              
              
                Your logit output shape is missing the class dimension. 
In my code snippet I’m creating the logits as [batch_size, nb_classes, height, width] and the target es [batch_size, height, width]. If you stick to these shapes, it should work. 

nn.CrossentropyLoss expects logits and uses F.log_softmax + nn…