diff --git a/unsloth/kernels/flex_attention.py b/unsloth/kernels/flex_attention.py
index 99e46a904c..1342bfafc2 100644
--- a/unsloth/kernels/flex_attention.py
+++ b/unsloth/kernels/flex_attention.py
@@ -42,15 +42,13 @@ if not HAS_FLEX_ATTENTION:
 
     # Below fails on compiled_autograd, so disable it
     try:
-        disable_compiled_autograd = torch._dynamo.compiled_autograd.disable
+        disable_compile = torch._dynamo.compiled_autograd.disable
     except:
-        disable_compiled_autograd = lambda f: f
+        disable_compile = lambda f: f
     pass
 
     # Logit softcapping
-    @disable_compiled_autograd(
-        torch.compile(fullgraph = True, dynamic = True, options = torch_compile_options)
-    )
+    @disable_compile(torch.compile(fullgraph = True, dynamic = True, options = torch_compile_options))
     def slow_attention_softcapping(Q, K, V, causal_mask, self, bsz, q_len):
         n_heads    = self.num_heads
         head_dim   = self.head_dim