Modifying dataset during training

Hello,

I need to modify data batches right after I used them to update my models, is there any existing way to do it when operating standard DataLoaders?

what I basically want is the last line in the following “code”:

trainset = TensorDataset(full_data, full_labels)
data_loader = DataLoader(trainset, batch_size=100, sampler = SubsetRandomSampler(indices))
for idx, (X, Y) in enumerate(data_loader):
    training step on (X, Y)
    **update: X  <- transform(X) directly in the trainset** 

Where transform() thing is also trained inside the loop, so I want to apply it and modify the dataset right after it was updated on X.

Any ideas would be extremely appreciated! Thank you in advance!