ptrblck:

However, we could just pass the BatchSampler as the sampler argument to the DataLoader (which I’ve done), so that we will get the complete batch of indices in __getitem__ .


Yes, that was exactly the case and your example worked perfectly. The only difference with my initial code was that I passed my CustomBatchSampler instance as batch_sampler, not sampler, but as I said (thanks to your example) is fixed now.



 ptrblck:

If your data is already in the system memory and you could just slice it to create batches, you most likely won’t see a speedup (and might even slow down the code due to some overhead in the DataLoader ).


Well, that explains the timing I got. For the record, the times of my runs were almost the same iterating over both objects (DataLoader or CustomBatchSampler). The difference was visible when I used num_workers > 0 and added torch.cuda.synchronize at the end of each loop, which confirms that the DataLoader workers fetch the data while the GPU is busy.
Thanks a lot @ptrblck!