meg:

For now I am doing a crossentropy loss computation but I don’t think it is the correct method.


I think it might be the correct method.
Based on your description it seems you are working on a multi-class segmentation use case, where each pixel belongs to one specific class.
Your model output shape and target shape look good, if you are dealing with 3 classes. The target should contain the class indices in the range [0, nb_classes-1] and you should be fine using nn.CrossEntropyLoss.
Let me know, if that’s your use case.