diff --git a/py/nodes/sageAttention.py b/py/nodes/sageAttention.py
index 6b24708..753fa34 100644
--- a/py/nodes/sageAttention.py
+++ b/py/nodes/sageAttention.py
@@ -27,7 +27,7 @@ def attention_sage(  # noqa: PLR0917
     attn_precision=None,
     skip_reshape=False,
     sageattn_allow_head_sizes: None | set | tuple | list = None,
-    sageattn_function=sageattention.sageattn,
+    sageattn_function=sageattention.sageattn if sageattention is not None else None,
     sageattn_verbose=False,
     **kwargs: dict[str],
 ):