Thanks for following up on this issue.
I don’t think the general workflow changes in the distributed setup, but think it’s limited by the way multiprocessing works.
The issue I’m seeing is that you need to manipulate the indices passed to Dataset.__getitem__ in each iteration.
The usual workflow would be to create the Dataset (with a custom sampler), setup the DataLoader, and iterate it for a complete epoch. If you are using multiprocessing, each worker will create a copy of the dataset, so that manipulating the underlying loader.dataset through the DataLoader won’t reflect the changes until the next epoch.
I’m not 100% sure about the behavior of the sampler, but would also assume it’s copied using num_workers>0.
Maybe a valid approach would be to use a shared array (e.g. via mp.Array as seen in this example 4), use it in the custom sampler, and apply the changes on the fly. However, now the distributed setup might be critical and you might need to use a locking mechanism to synchronize the manipulation of this array.
I hope someone else has a better idea how this use case can be properly implemented.