fix: remove the bug in the graph cleaning function

Co-authored-by: Jesper Dramsch <[email protected]>
ecmwf · Aug 14, 2024 · 6997711 · 6997711
1 parent 3d8b1ea
commit 6997711
Showing 1 changed file with 5 additions and 3 deletions.
diff --git a/src/anemoi/graphs/create.py b/src/anemoi/graphs/create.py
@@ -61,10 +61,12 @@ def clean(self, graph: HeteroData) -> HeteroData:
         HeteroData
             cleaned graph
         """
+        LOGGER.info("Cleaning graph.")
         for type_name in chain(graph.node_types, graph.edge_types):
-            for attr_name in list(graph[type_name].keys()):
-                if attr_name.startswith("_"):
-                    del graph[type_name][attr_name]
+            attr_names_to_remove = [attr_name for attr_name in graph[type_name] if attr_name.startswith("_")]
+            for attr_name in attr_names_to_remove:
+                del graph[type_name][attr_name]
+                LOGGER.info(f"{attr_name} deleted from graph.")
 
         return graph