@faizan_shaikh. Sorry for late response.
This is my codesnipet for dataloader.
class TrainDataset(Dataset):
    def __init__(self, root="", live_sub_dirs=[], fake_sub_dirs=[]):
        self.live_sub_dirs = live_sub_dirs
        self.fake_sub_dirs = fake_sub_dirs

        self.pos_filelist = get_total_list(root, live_sub_dirs)
        self.neg_filelist = get_total_list(root, fake_sub_dirs)

        count = len(self.neg_filelist) // len(self.pos_filelist)
        self.total_filelist = np.asarray(self.pos_filelist * count + self.neg_filelist, dtype=np.str)
        np.random.shuffle(self.total_filelist)

        self.transform = transforms.Compose([
            transforms.RandomCrop(128),
            transforms.RandomHorizontalFlip(),
            transforms.ToTensor()
        ])
        self.aug = strong_aug(0.5)

    def __getitem__(self, idx):
        img_path = self.total_filelist[idx]
        img = Image.open(img_path).convert("RGB")
        img = self.aug(image=np.array(img))["image"]
        img = self.transform(Image.fromarray(img))

        target = 0 if str('fake') in img_path.lower() else 1

        return img, torch.tensor(target, dtype=torch.long)
    
    def __len__(self):
        return len(self.total_filelist)