Hi @ptrblck, if I understood you correctly something like the following MWE is what you’re describing?
class MyDataset(Dataset):
     def __init__(self, *args, epoch=0):
         self.epoch = epoch

     def __getittem__(self, index):
         x, y = self.data[index]
         if self.epoch == epoch:
            x = transforms(x)
         return x, y

# then during training
for epoch in range(epochs):
     dataset.epoch = epoch
     for x, y in loader:
          # train as usual