diff --git a/src/penai/hierarchy_generation/inference.py b/src/penai/hierarchy_generation/inference.py
index ebdc347..43a4dba 100644
--- a/src/penai/hierarchy_generation/inference.py
+++ b/src/penai/hierarchy_generation/inference.py
@@ -72,6 +72,8 @@ def from_penpot_shape(cls, shape: PenpotShapeElement) -> Self:
         for child in shape.get_direct_children_shapes():
             element.children.append(cls.from_penpot_shape(child))
 
+        return element
+
     def flatten(self) -> Iterable[Self]:
         yield self