Thanks @ptrblck. We use 16 workers and the current dataloader is
 train_transform = transforms.Compose([
           transforms.RandomHorizontalFlip(),
           transforms.ToTensor(),
           transforms.Normalize(mean = RGB_MEAN,
                            std = RGB_STD),])
    dataset_train = FaceDataset(DATA_ROOT, RECORD_DIR, train_transform)
    train_sampler = torch.utils.data.distributed.DistributedSampler(dataset_train)
    train_loader = torch.utils.data.DataLoader(dataset_train, batch_size=batch_size, shuffle = (train_sampler is None), num_workers=workers, pin_memory=True, sampler=train_sampler, drop_last=True)
    SAMPLE_NUMS = dataset_train.get_sample_num_of_each_class()
    NUM_CLASS = len(train_loader.dataset.classes)
 

And the current dataloader is refered from https://github.com/HuangYG123/CurricularFace/blob/8b2f47318117995aa05490c05b455b113489917e/dataset/datasets.py#L92
path, target = self.imgs[index]
sample = Image.open(path)
sample = sample.convert("RGB")
if self.transform is not None:
     sample = self.transform(sample)

Is it fine to replace your customized transformation for the above code?