meg:

but I actually have 4 classes


In that case the output should have the shape [batch_size, 4, height, width], where each channel corresponds to the logit for a class.



 meg:

Moreover, when I am loading my label I first used the transform ToTensor() (but it puts everything in the range from 0 to 1) so I changed it to a transform that replaces it with integers form [0, 3] (because there are 4 classes). Is it correct to do it that way ?


Some transformations on the mask are dangerous, as they might corrupt the mask.
ToTensor is such an operation and you should make sure that your mask contains the class indices in the range [0, 3], as you described.
Also, if you are reshaping, I would recommend to use the NEAREST interpolation method, as others might also corrupt the class indices.