diff --git a/src/anemoi/graphs/edges/builder.py b/src/anemoi/graphs/edges/builder.py
index 1b22b2e..d502bee 100644
--- a/src/anemoi/graphs/edges/builder.py
+++ b/src/anemoi/graphs/edges/builder.py
@@ -119,10 +119,8 @@ def update_graph(self, graph: HeteroData, attrs_config: DotDict | None = None) -
         """
         graph = self.register_edges(graph)
 
-        if attrs_config is None:
-            return graph
-
-        graph = self.register_attributes(graph, attrs_config)
+        if attrs_config is not None:
+            graph = self.register_attributes(graph, attrs_config)
 
         return graph