diff --git a/unsloth/kernels/flex_attention.py b/unsloth/kernels/flex_attention.py
index 678574928c..99e46a904c 100644
--- a/unsloth/kernels/flex_attention.py
+++ b/unsloth/kernels/flex_attention.py
@@ -44,7 +44,7 @@ if not HAS_FLEX_ATTENTION:
     try:
         disable_compiled_autograd = torch._dynamo.compiled_autograd.disable
     except:
-        disable_compiled_autograd = lambda *args, **kwargs: *args, **kwargs
+        disable_compiled_autograd = lambda f: f
     pass
 
     # Logit softcapping