Last I looked, the dataloader would call the sampler at the beginning of an epoch. Given that instantiating a dataloader isn’t that much more expensive than starting an epoch (i.e. instantiating the iterator), I’m not sure that what you have in mind is easy to implement or much more efficient.
The obvious other solution is to forego the dataloading and produce a dataset (possibly an iterative dataset) that accomplishes this (in a well-optimized pipeline having multiple workers might not be as crucial as it seems to be initially).
That said, while I do believe that using the sampler is likely a good idea, it’s totally possible that you’ll have a much better solution matching your workflow. And if that’s the case, I’d be really keen to hear how you did it.
Best regards
Thomas