Data augmentation with transform after create dataloader


I’m trying to automate a training pipeline for my project with pytorch and sklearn cross-validation.
I know I can do transformations while creating the dataset, but in the pipeline I first concatenate all data to split with the cross-validation method. So, if I use the transform in the dataset creation the transformations will be applied to training and testing data.

I would like to know if it is possible to make the transformations in the data after creating the dataloader. This way I can separate the transforms for train and test.