torch dataloader crashed after `set_up_backend` about torchquad

[gomezzz]

Hi @jk4011 !

Thanks for the bug report. So, my first suspicion would be that your data generator returns data that is on the CPU while you set up the default device to be GPU with set_up_backend("torch")?

Could you by any chance post a minimal code example to help me reproduce the problem? Then I would have a more detailed look.

Thanks!

[jk4011]

Hello @gomezzz
Of course!

import torch
from torchquad import set_up_backend
from torchvision import datasets
from torch.utils.data import DataLoader

# this error only occurs when cuda is available
assert(torch.cuda.is_available())

set_up_backend('torch')

training_data = datasets.MNIST(
    root="data",
    train=True,
    download=True,
)

train_loader = DataLoader(training_data, batch_size=64, shuffle=True)

for (i, x) in enumerate(train_loader):
    pass

and the error is like below

Traceback (most recent call last):
  File "/home/wlsgur4011/torchlevy/tmp.py", line 19, in <module>
    for (i, x) in enumerate(train_loader):
  File "/home/wlsgur4011/.anaconda3/envs/torchlevy/lib/python3.8/site-packages/torch/utils/data/dataloader.py", line 681, in __next__
    data = self._next_data()
  File "/home/wlsgur4011/.anaconda3/envs/torchlevy/lib/python3.8/site-packages/torch/utils/data/dataloader.py", line 720, in _next_data
    index = self._next_index()  # may raise StopIteration
  File "/home/wlsgur4011/.anaconda3/envs/torchlevy/lib/python3.8/site-packages/torch/utils/data/dataloader.py", line 671, in _next_index
    return next(self._sampler_iter)  # may raise StopIteration
  File "/home/wlsgur4011/.anaconda3/envs/torchlevy/lib/python3.8/site-packages/torch/utils/data/sampler.py", line 247, in __iter__
    for idx in self.sampler:
  File "/home/wlsgur4011/.anaconda3/envs/torchlevy/lib/python3.8/site-packages/torch/utils/data/sampler.py", line 132, in __iter__
    yield from torch.randperm(n, generator=generator).tolist()
RuntimeError: Expected a 'cuda' device type for generator but found 'cpu'

[gomezzz]

Hi @jk4011

Sorry for the late reply 🙏 Busy days

Just tried the snippet and indeed same happens for me. However, I think, the problem is due to a missing device specification.

See also this stackoverflow post

Below snippet fixes the problem (and a follow-up one related to MNIST being PIL images) for me:

import torch
from torchquad import set_up_backend
from torchvision import datasets
from torch.utils.data import DataLoader
from torchvision import transforms

# this error only occurs when cuda is available
assert torch.cuda.is_available()

set_up_backend("torch")

transform = transforms.Compose(
    [
        # you can add other transformations in this list
        transforms.ToTensor()
    ]
)

training_data = datasets.MNIST(
    root="data", train=True, download=True, transform=transform
)

train_loader = DataLoader(
    training_data,
    batch_size=64,
    shuffle=True,
    generator=torch.Generator(device="cuda"),
)

for (i, x) in enumerate(train_loader):
    pass

Does it fix it for you too, @jk4011 ? :)

[gomezzz]

Should be fixed / not a problem in torchquad