Shape of ground truth in multiclass image segmentation with pytorch

Question

I\'m working on 128 x 128 x 3 cell images and want to segment them into 5 classes including backgrounds. I first made target images to be 128 x 128 and values are in {0,1,2,3,4}