Knn_interpolate inside a forward() function of main model class causes RuntimeError

username: pserg
cooked: <p>Hello, I’m trying to build a model using pytorch_geometric and my main forward() function calls<br>
pytorch_geometric.nn.unpool.knn_interpolate function.</p>
<p>And I get this RuntimeError:</p>
<pre><code class="lang-auto">&lt;ipython-input-18-bbb21c6c9666&gt; in train(params, simulator, train_loader, valid_loader, valid_rollout_dataset)
     21             optimizer.zero_grad()
     22             data = data.cuda()
---&gt; 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):
-&gt; 1501             return forward_call(*args, **kwargs)
   1502         # Do not call functions when jit is used
   1503         full_backward_hooks, non_full_backward_hooks = [], []

&lt;ipython-input-24-cadb18fc6e40&gt; in forward(self, data)
     76         # perform 2nd unpooling
     77         with torch.no_grad():
---&gt; 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():
---&gt; 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     """
---&gt; 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,
           ~~~~~~~~~~~~~~~~~~~~~~~~~~~ &lt;--- 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
</code></pre>
<p>I’m using PyTorch has version 2.0.0+cu118 with cuda 11.8 on Google Colab:</p>
<pre><code class="lang-auto">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
</code></pre>
<p>Also please note, the RuntimeError happens at the time of the second call. (I’m calling this function 2 times consecutively).</p>