Hello, I’m trying to build a model using pytorch_geometric and my main forward() function calls
pytorch_geometric.nn.unpool.knn_interpolate function.
And I get this RuntimeError:
<ipython-input-18-bbb21c6c9666> in train(params, simulator, train_loader, valid_loader, valid_rollout_dataset)
21 optimizer.zero_grad()
22 data = data.cuda()
---> 23 pred = simulator(data)
24 loss = loss_fn(pred, data.y)
25 loss.backward()
/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py in _call_impl(self, *args, **kwargs)
1499 or _global_backward_pre_hooks or _global_backward_hooks
1500 or _global_forward_hooks or _global_forward_pre_hooks):
-> 1501 return forward_call(*args, **kwargs)
1502 # Do not call functions when jit is used
1503 full_backward_hooks, non_full_backward_hooks = [], []
<ipython-input-24-cadb18fc6e40> in forward(self, data)
76 # perform 2nd unpooling
77 with torch.no_grad():
---> 78 res = pyg.nn.unpool.knn_interpolate(x = node_features,
79 pos_x = graph.position_seq[:, -1], pos_y = latent_graph.position_seq[:,-1], k=5) # changed: target_position = position_seq[:, -1]
80
/usr/local/lib/python3.10/dist-packages/torch_geometric/nn/unpool/knn_interpolate.py in knn_interpolate(x, pos_x, pos_y, batch_x, batch_y, k, num_workers)
45
46 with torch.no_grad():
---> 47 assign_index = knn(pos_x, pos_y, k, batch_x=batch_x, batch_y=batch_y,
48 num_workers=num_workers)
49 y_idx, x_idx = assign_index[0], assign_index[1]
/usr/local/lib/python3.10/dist-packages/torch_geometric/nn/pool/__init__.py in knn(x, y, k, batch_x, batch_y, cosine, num_workers)
90 :rtype: :class:`torch.Tensor`
91 """
---> 92 return torch_cluster.knn(x, y, k, batch_x, batch_y, cosine, num_workers)
93
94
RuntimeError: The following operation failed in the TorchScript interpreter.
Traceback of TorchScript (most recent call last):
File "/usr/local/lib/python3.10/dist-packages/torch_cluster/knn.py", line 72, in knn
ptr_y = torch.bucketize(arange, batch_y)
return torch.ops.torch_cluster.knn(x, y, ptr_x, ptr_y, k, cosine,
~~~~~~~~~~~~~~~~~~~~~~~~~~~ <--- HERE
num_workers)
RuntimeError: y.device().is_cpu() INTERNAL ASSERT FAILED at "csrc/cpu/knn_cpu.cpp":14, please report a bug to PyTorch. y must be CPU tensor
I’m using PyTorch has version 2.0.0+cu118 with cuda 11.8 on Google Colab:
DISTRIB_ID=Ubuntu
DISTRIB_RELEASE=20.04
DISTRIB_CODENAME=focal
DISTRIB_DESCRIPTION="Ubuntu 20.04.5 LTS"
NAME="Ubuntu"
VERSION="20.04.5 LTS (Focal Fossa)"
ID=ubuntu
ID_LIKE=debian
PRETTY_NAME="Ubuntu 20.04.5 LTS"
VERSION_ID="20.04"
HOME_URL="https://www.ubuntu.com/"
SUPPORT_URL="https://help.ubuntu.com/"
BUG_REPORT_URL="https://bugs.launchpad.net/ubuntu/"
PRIVACY_POLICY_URL="https://www.ubuntu.com/legal/terms-and-policies/privacy-policy"
VERSION_CODENAME=focal
UBUNTU_CODENAME=focal
Also please note, the RuntimeError happens at the time of the second call. (I’m calling this function 2 times consecutively).