Hi, @vmirly1 Thanks for your response.
I have tried the code you offered as follows :
    dst = MyDataset(root = args.dataset_root_dir, set_name="train")
    #print(len(dst)) 
    data_loader = DataLoader(dataset=dst, batch_size=args.batch_size, shuffle=False)
    batch_x, batch_y = next(iter(data_loader))
    print(batch_x.shape, batch_y.shape)

And the shape of data and target like this:

torch.Size([10, 3, 512, 512])   torch.Size([10, 1, 512, 512])

It confirms that MyDataset has no problem.
I found that  the dimension mismatch error occurs when I do TF.vflip() and TF.hflip() transforms, and I solve the problem by this:
def transform(self, image, mask):
        # Resize
        resize = transforms.Resize(size=(args.input_size, args.input_size))
        image = resize(image)
        mask = resize(mask)
        # Random horizontal flipping
        if random.random() > 0.5:
            image = TF.hflip(image)
            mask = TF.hflip(image)
        # Random vertical flipping
        if random.random() > 0.5:
            image = TF.vflip(image)
            mask = TF.vflip(mask)
        mask = TF.to_grayscale(mask)
        # Transform to Tensor
        image = TF.to_tensor(image)
        mask = TF.to_tensor(mask)
        image = TF.normalize(image, mean=[0.4353, 0.4452, 0.4131],
                                    std=[0.2044, 0.1924, 0.2013])
        return image, mask

It means I should read the images  and masks  in RGB and after transfroms the masks then convert to grayscale, and then I can get masks [batch_size, 1, height, width] in DataLoader. Is there better solutions?
Oh, I will go to re-edit the topic name. Thanks a lot