fix coalesce on CUDA

pyg-team · Oct 26, 2021 · 101ed2c · 101ed2c
1 parent 2910199
commit 101ed2c
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/torch_geometric/utils/coalesce.py b/torch_geometric/utils/coalesce.py
@@ -65,7 +65,8 @@ def coalesce(
         return edge_index
 
     dim_size = edge_index.size(1)
-    idx = torch.arange(0, nnz).sub_(mask.logical_not_().cumsum(dim=0))
+    idx = torch.arange(0, nnz, device=edge_index.device)
+    idx.sub_(mask.logical_not_().cumsum(dim=0))
 
     if isinstance(edge_attr, Tensor):
         edge_attr = scatter(edge_attr, idx, 0, None, dim_size, reduce)