Thanks again, @ptrblck! I also had to use batch_size=None to prevent the collate function to add an additional dimension. So the final snippet is
import torch
import torch.utils.data as data

class MyDataset(data.Dataset):
    def __init__(self):
        self.data = torch.rand(100, 5)

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

    def __getitem__(self, i) -> torch.Tensor:
        print(i)
        return self.data[i]

dataset = MyDataset()
sampler = data.BatchSampler(data.RandomSampler(dataset), 2, False)
loader  = data.DataLoader(dataset, sampler=sampler, batch_size=None)

for x in loader:
    print(x)

By the way, I feel like using the keyword sampler instead batch_sampler for specifically sampling batches is a bit weird. Perhaps, the interface could be improved a little ? Although, if it ain’t broke, don’t fix it