This is how i usedyour code:
I dont have that error anymore, but is the proper way to use your code?
“This is the current error”
“ValueError: Expected target size (1, 480), got torch.Size([1, 3, 480])”
“Used nn.CrossEntropyLoss(), num_classes=256(randomly initialized, bcz i was getting the len(colors) as 2992”
class Cus_dataset(torch.utils.data.Dataset):
def __init__(self, root,transform,transformm):

    self.root = root

    self.transform =  transform

    self.transformm = transformm

    # load all image files, sorting them to

    # ensure that they are aligned

    self.imgs = list(sorted(os.listdir(os.path.join(root, "original_images"))))

    self.masks = list(sorted(os.listdir(os.path.join(root, "col"))))

                  

def __getitem__(self, idx):

    # load images ad masks

    img_path = os.path.join(self.root, "original_images", self.imgs[idx])

    mask_path = os.path.join(self.root, "col", self.masks[idx])

    img = Image.open(img_path).convert("RGB")

    # note that we haven't converted the mask to RGB,

    # because each color corresponds to a different instance

    # with 0 being background

                    

    mask = Image.open(mask_path)

    mask = np.array(mask)

    #colors = np.unique(mask)

    #mask =self.transform(mask)

    def mask_to_class(mask):

      target = torch.from_numpy(mask)

      h,w = target.shape[0],target.shape[1]

      masks = torch.empty(h, w, dtype=torch.long)

      colors = torch.unique(target.view(-1,target.size(2)),dim=0).numpy()

      target = target.permute(2, 0, 1).contiguous()

      mapping = {tuple(c): t for c, t in zip(colors.tolist(), range(len(colors)))}

      for k in mapping:

        idx = (target==torch.tensor(k, dtype=torch.uint8).unsqueeze(1).unsqueeze(2))

        validx = (idx.sum(0) == 3) 

        masks[validx] = torch.tensor(mapping[k], dtype=torch.long)

      

    masks = mask_to_class(mask)
   return masks