diff --git a/flatten.py b/flatten.py
index c180564..574c118 100644
--- a/flatten.py
+++ b/flatten.py
@@ -6,8 +6,22 @@ from specklepy.objects import Base
 
 
 def flatten_base(base: Base) -> Iterable[Base]:
-    """Take a base and flatten it to an iterable of bases."""
-    if hasattr(base, "elements") and base.elements is not None:
-        for element in base["elements"]:
+    """Flatten a base object into an iterable of bases.
+    
+    This function recursively traverses the `elements` or `@elements` attribute of the 
+    base object, yielding each nested base object.
+
+    Args:
+        base (Base): The base object to flatten.
+
+    Yields:
+        Base: Each nested base object in the hierarchy.
+    """
+    # Attempt to get the elements attribute, fallback to @elements if necessary
+    elements = getattr(base, "elements", getattr(base, "@elements", None))
+    
+    if elements is not None:
+        for element in elements:
             yield from flatten_base(element)
+    
     yield base