fix

pyg-team · Oct 25, 2022 · d3b278a · d3b278a
1 parent 3b5edb3
commit d3b278a
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/torch_geometric/transforms/virtual_node.py b/torch_geometric/transforms/virtual_node.py
@@ -50,7 +50,7 @@ def __call__(self, data: Data) -> Data:
                 if key == 'edge_weight':
                     size[dim] = 2 * num_nodes
                     fill_value = 1.
-                if key == 'batch':
+                elif key == 'batch':
                     size[dim] = 1
                     fill_value = int(value[0])
                 elif old_data.is_edge_attr(key):