tnt has SplitDataset 223.
It has method to select the partition and therefore, changes the data from __getitem__.
I’m not too sure how it interacts with DataLoader though. DataLoader seems to make an index or something before returning the iterator. So modifying the dataset by selecting partition might break DataLoader.