If you are using nn.CrossEntropyLoss, you should have to use this approach:



 ptrblck:

nn.CrossEntropyLoss expects the model output to have the shape [batch_size, nb_classes, height, width] and the target [batch_size, height, width] containing the class indices in the range [0, nb_classes-1] .


Your target mask should not contain a channels dimension and for 5 classes should contain the class indices in the range [0, 4]. If the background class is not included in these 5 classes, then you should use 6. It doesn’t matter, if the background is class index 0 or any other.