By default, data.DataLoader indexes elements of a batch one by one and collates them back into tensors. I have a dataset (subclass of data.Dataset) which can be indexed (efficiently) by slices. For example, the following can be indexed by slices:
class MyDataset(data.Dataset):
   def __init__(self, a: Tensor, b: Tensor):
      self.a, self.b = a, b

   def __len__(self) -> int:
        return len(self.a)

   def __getitem__(self, i) -> Tuple[Tensor, Tensor]:
       return self.a[i], self.b[i]

I would like the DataLoader to take advantage of that. I think I need to use a data.Sampler, but I didn’t get at all how to from the documentation. Any idea ?