bs = 2
num_epochs = 100
learning_rate = 1e-3
mom  = 0.9
class SEMDataTrain(data.Dataset):

    def __init__(self, image_path, mask_path, in_size=512, out_size=512 ,transforms = None):
        """
        Args:
            image_path (str): the path where the image is located
            mask_path (str): the path where the mask is located
            option (str): decide which dataset to import
        """
        # all file names
        self.mask_path = mask_path
        self.image_path = image_path
        self.mask_arr = os.listdir(mask_path)
        self.image_arr = os.listdir(image_path)
        self.in_size, self.out_size = in_size, out_size
        # Calculate len
        self.data_len = len(self.mask_arr)
    
    def transform(self, img_as_img, msk_as_img):
        # Resize
        resize = transforms.Resize(size=(256, 256))
        msk_as_img = resize(img_as_img)
        msk_as_img = resize(msk_as_img)
         
    
    def __getitem__(self,idx):
        """Get specific data corresponding to the index
        Args:
            index (int): index of the data
        Returns:
            Tensor: specific data on index which is converted to Tensor
        """
        """
        # GET IMAGE
        """
        single_image_name = self.image_arr[idx]
        img_as_img = Image.open(os.path.join(self.image_path,single_image_name))
        # img_as_img.show()
        img_as_np = np.asarray(img_as_img)
        mg_as_img = Image.fromarray(img_as_np)
        img_as_img = img_as_img.convert('L')
        plt.imshow(img_as_img)
        plt.show()
        
        # Augmentation
        # flip {0: vertical, 1: horizontal, 2: both, 3: none}
        flip_num = randint(0, 3)
        img_as_np = flip(img_as_np, flip_num)

        # Noise Determine {0: Gaussian_noise, 1: uniform_noise
        if randint(0, 1):
            # Gaussian_noise
            gaus_sd, gaus_mean = randint(0, 20), 0
            img_as_np = add_gaussian_noise(img_as_np, gaus_mean, gaus_sd)
        else:
            # uniform_noise
            l_bound, u_bound = randint(-20, 0), randint(0, 20)
            img_as_np = add_uniform_noise(img_as_np, l_bound, u_bound)

        # Brightness
        pix_add = randint(-20, 20)
        img_as_np = change_brightness(img_as_np, pix_add)

        # Elastic distort {0: distort, 1:no distort}
        sigma = randint(6, 12)
        # sigma = 4, alpha = 34
        img_as_np, seed = add_elastic_transform(img_as_np, alpha=34, sigma=sigma, pad_size=20)

        # Crop the image
        img_height, img_width = img_as_np.shape[0], img_as_np.shape[1]
        pad_size = int((self.in_size - self.out_size)/2)
        img_as_np = np.pad(img_as_np, pad_size, mode="symmetric")
        y_loc, x_loc = randint(0, img_height-self.out_size), randint(0, img_width-self.out_size)
        img_as_np = cropping(img_as_np, crop_size=self.in_size, dim1=y_loc, dim2=x_loc)
        '''
        # Sanity Check for image
        img1 = Image.fromarray(img_as_np)
        img1.show()
        '''
        # Normalize the image
        img_as_np = normalization2(img_as_np, max=1, min=0)
        img_as_np = np.expand_dims(img_as_np, axis=0)  # add additional dimension
        img_as_tensor = torch.from_numpy(img_as_np).float()  # Convert numpy array to tensor
        
        

        """
        # GET MASK
        """
        single_mask_name = self.mask_arr[idx]
        msk_as_img = Image.open(os.path.join(self.mask_path,single_mask_name))
        # msk_as_img.show()
        msk_as_np = np.asarray(msk_as_img)

        # flip the mask with respect to image
        msk_as_np = flip(msk_as_np, flip_num)

        # elastic_transform of mask with respect to image

        # sigma = 4, alpha = 34, seed = from image transformation
        msk_as_np, _ = add_elastic_transform(
            msk_as_np, alpha=34, sigma=sigma, seed=seed, pad_size=20)
        msk_as_np = approximate_image(msk_as_np)
        

        # Crop the mask
        msk_as_np = cropping(msk_as_np, crop_size=self.out_size, dim1=y_loc, dim2=x_loc)
        '''
         Sanity Check for mask
        img2 = Image.fromarray(msk_as_np)
        img2.show()
        '''

        # Normalize mask to only 0 and 1
        msk_as_np = msk_as_np/255
        # msk_as_np = np.expand_dims(msk_as_np, axis=0)  # add additional dimension
        msk_as_tensor = torch.from_numpy(msk_as_np).long()  # Convert numpy array to tensor
        

        return (img_as_tensor, msk_as_tensor)
    
         
    def __len__(self):
        """
        Returns:
            length (int): length of the data
            
        """
        return self.data_len

yes it works if with RGB images but sicnce the labels are gray scale it is not work