bs = 2
num_epochs = 100
learning_rate = 1e-3
mom  = 0.9
class MYDataLoader(data.Dataset):
    def __init__(self,root_dir,seg_dir,transforms = None):
        self.root_dir = root_dir
        self.seg_dir = seg_dir
        self.transforms = transforms
        self.files = os.listdir(self.root_dir)
        self.lables = os.listdir(self.seg_dir)
    
    def __len__(self):
        return len(self.files)
    
    def __getitem__(self,idx):
        img_name = self.files[idx]
        label_name = self.lables[idx]
        img = Image.open(os.path.join(self.root_dir,img_name))
        label = Image.open(os.path.join(self.seg_dir,label_name))
        if self.transforms:
            img = self.transforms(img)
            label = self.transforms(label)
            return img,label
        else:
            return img, label
full_dataset = MYDataLoader('/ data/training/images',
                                     '//data/training/labels',
                                    transforms=tfms.Compose([tfms.Resize((256,256)),
    tfms.ColorJitter(hue=.05, saturation=.05),
    tfms.RandomHorizontalFlip(),tfms.Grayscale(num_output_channels=1),
    tfms.ToTensor() , tfms.Normalize([0.5],[0.5])
]))
    

this is the code with Normalize module.