Usually you would use nn.CrossEntropyLoss for a multi-class segmentation, which would therefore dictate the output shape as [batch_size, nb_classes, height, width] and the labels as [batch_size, height, width] containing the class indices.
Since you are using nn.MSELoss, could you explain the output shape of your model and how you achieved this shape?



 Phanikumar2all:

When I check the values of masks_train, it is showing all zero values.


Is this expected for this sample, i.e. does this sample only contain class0 pixels?
If not, then your mask creation might have some bugs.