Indices on the same device as the indexed tensor

Signed-off-by: ddelange <[email protected]>
ray-project · Jan 12, 2023 · b30b267 · b30b267
1 parent d5b7e15
commit b30b267
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/python/ray/train/examples/pytorch_geometric/distributed_sage_example.py b/python/ray/train/examples/pytorch_geometric/distributed_sage_example.py
@@ -45,7 +45,7 @@ def test(self, x_all, subgraph_loader):
             for batch_size, n_id, adj in subgraph_loader:
 
                 edge_index, _, size = adj
-                x = x_all[n_id]
+                x = x_all[n_id.to(x_all.device)].to(train.torch.get_device())
                 x_target = x[: size[1]]
                 x = self.convs[i]((x, x_target), edge_index)
                 if i != self.num_layers - 1: