diff --git a/src/pyplaml/diagram_layout.py b/src/pyplaml/diagram_layout.py
index 4d919b8657b31dca9ac050cacb26e0e992d23bb2..e0f1df0a1dff3e5725d9774f8ea6357fc419e1f3 100644
--- a/src/pyplaml/diagram_layout.py
+++ b/src/pyplaml/diagram_layout.py
@@ -18,9 +18,9 @@ class DiagramLayout(ABC):
     def get_graph(self) -> nx.DiGraph:
         g = nx.DiGraph()
         for name, o in self.objects.items():
-            g.add_node(o.name)
+            g.add_node(o.get_key())
             for e in o.edges:
-                g.add_edge(e.target.get_key(), o.name)
+                g.add_edge(e.target.get_key(), o.get_key())
         return g