Solved by ptrblck in post #2 
              
              
                You could use an internal attribute in your custom Dataset to tag the samples, which should be augmented, and use it in the __getitem__ as a condition. At the beginning you could leave it empty and manipulate it after each epoch via directly accessing the internal dataset through the DataLoader. 
Th…