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.



 Alex_Ge:

Also, would you recommend CrossEntropyLoss , NLLloss or some other function for this sort of thing, and why one over the other?


nn.CrossentropyLoss expects logits and uses F.log_softmax + nn.NLLLoss internally, so these approaches will yield the same result.