ValueError when using BalancedBatchSampler

[SergioQuijanoRey]

Hi! I am using your custom sampler for a project of mine. I am currently using MNIST dataset loaded from torchvision. When I use your sampler I got the following error:

ValueError: Caught ValueError in DataLoader worker process 0.
Original Traceback (most recent call last):
  File "/usr/local/lib/python3.7/dist-packages/torch/utils/data/_utils/worker.py", line 287, in _worker_loop
    data = fetcher.fetch(index)
  File "/usr/local/lib/python3.7/dist-packages/torch/utils/data/_utils/fetch.py", line 49, in fetch
    data = [self.dataset[idx] for idx in possibly_batched_index]
  File "/usr/local/lib/python3.7/dist-packages/torch/utils/data/_utils/fetch.py", line 49, in <listcomp>
    data = [self.dataset[idx] for idx in possibly_batched_index]
  File "/usr/local/lib/python3.7/dist-packages/torchvision/datasets/mnist.py", line 127, in __getitem__
    img, target = self.data[index], int(self.targets[index])
ValueError: only one element tensors can be converted to Python scalars

I don't know what may be causing this error. Maybe different pytorch versions? I am currently using 1.10.0. Thanks in advance for your time 😄