diff --git a/main.py b/main.py
index e1bacd6..b44923e 100644
--- a/main.py
+++ b/main.py
@@ -9,6 +9,7 @@ from speckle_automate import (
     AutomationContext,
     execute_automate_function,
 )
+from specklepy.objects import Base
 
 from flatten import flatten_base