Why EdgeSamplerInput moves data to cpu?

[johnarevalo]

When torch.set_default_device('cuda:0'), iterating a LinkNeighborLoader object raises:

torch/utils/data/sampler.py", line 164, in __iter__
    yield from torch.randperm(n, generator=generator).tolist()

RuntimeError: Expected a 'cuda' device type for generator but found 'cpu'

If I add generator=torch.Generator(device='cuda:0') to the LinkNeighborLoader initialization, the error changes to:

torch_geometric/loader/link_loader.py", line 209, in collate_fn
     input_data: EdgeSamplerInput = self.input_data[index]
                                    ~~~~~~~~~~~~~~~^^^^^^^
torch_geometric/sampler/base.py", line 154, in __getitem__
     self.row[index],

RuntimeError: indices should be either on cpu or on the same device as the indexed tensor (cpu)

in 154 line index tensor is in cuda:0 but self.row is in cpu. It seems that is forced in:

pytorch_geometric/torch_geometric/sampler/base.py

Lines 133 to 139 in 36bc925

	input_id = input_id.cpu()
	row = row.clone().cpu()
	col = col.clone().cpu()
	if label is not None:
	label = label.cpu()
	if time is not None:
	time = time.cpu()

What is the reason to enforce this?

[rusty1s]

Thanks for reporting. It looks like loaders are a bit fragile if you specify CUDA as the default device. I will look into this.