mask

torch_cluster/radius.py

@@ -69,9 +69,8 @@ def radius(x, y, r, batch_x=None, batch_y=None, max_num_neighbors=32):
     col = [torch.tensor(c) for c in col]
     row = [torch.full_like(c, i) for i, c in enumerate(col)]
     row, col = torch.cat(row, dim=0), torch.cat(col, dim=0)
-    row = row[col<tree.n]
-    col = col[col<tree.n]
-    return torch.stack([row, col], dim=0)
+    mask = col < tree.n
+    return torch.stack([row[mask], col[mask]], dim=0)
 
 
 def radius_graph(x, r, batch=None, loop=False, max_num_neighbors=32):