Rasha_Salim:

That’s why I added this line in the Dataset class : mask = mask[0, :, :]


I understand that it is a 6-class segmentation problem?
I am not sure what the initial target shape [16, 6, 224, 224] means. Is it [batch size, num classes (one hot), H, W]?
If the target was one hot, you should do mask = torch.argmax(mask, dim=1).



 Rasha_Salim:

here where things get confusing for me, because the output of me model is torch.Size([10, 6, 224, 224]).


This is confusing now. If the batch size is 16, the output of the model should have been [16, 6,224,224].