I’m not aware of the transforms.PILToTensor or transforms.Tensor methods. Do you mean transforms.ToTensor? If so, what are you comparing it against?


That’s incorrect for a multi-class segmentation, where each pixel belongs to one specific class only. In that case the target tensor should have the shape [batch_size, height, width] and contain the class indices in the range [0, nb_classes-1]. Since your tensor seems to be a one-hot encoded mask, you could simply call target = one_hot_target.argmax(2).


For a multi-class segmentation you should use nn.CrossEntropyLoss, remove the last sigmoid and pass the raw logits in the shape [batch_size, nb_classes, height, width] to the criterion.