Dataloader with two sets of augmentation on same image

I am trying to do weak and strong augmentation of the same set of images by maintaining the actual correspondence. I tried with ConcatDataset. Though images come from two sets of augmentations, it doesn’t maintain the one-to-one correspondence in the first half and 2nd half of the batch.
Please let me know if you have any idea.
Thank You

ConcatDataset will concatenate multiple Datasets sequentially, i.e. the len of the ConcatDataset will be the sum of all passed Datasets.
If you want to augment the same sample in a different way, you could use different transformations in the __getitem__ of the Dataset and return both transformed samples.
This would double your batch size, so that you could lower it in the DataLoader, if needed.