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)