Take a look at this post to see how a BatchSampler can be used to provide a all indices of the current batch to the Dataset.__getitem__ and write a custom implementation sampling from your datasets by reusing the BatchSampler implementation.
1 Like