improfe has_isolated performance

pyg-team · Feb 25, 2022 · 970de13 · 970de13
1 parent 1cbf696
commit 970de13
Showing 1 changed file with 2 additions and 3 deletions.
diff --git a/torch_geometric/utils/isolated.py b/torch_geometric/utils/isolated.py
@@ -17,9 +17,8 @@ def contains_isolated_nodes(edge_index, num_nodes=None):
     :rtype: bool
     """
     num_nodes = maybe_num_nodes(edge_index, num_nodes)
-    (row, col), _ = remove_self_loops(edge_index)
-
-    return torch.unique(torch.cat((row, col))).size(0) < num_nodes
+    edge_index, _ = remove_self_loops(edge_index)
+    return torch.unique(edge_index.view(-1)).size(0) < num_nodes
 
 
 def remove_isolated_nodes(edge_index, edge_attr=None, num_nodes=None):