Update flex_attention.py

unslothai · Sep 3, 2024 · 7ea6395 · 7ea6395
1 parent 315136a
commit 7ea6395
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/unsloth/kernels/flex_attention.py b/unsloth/kernels/flex_attention.py
@@ -47,7 +47,7 @@
 pass
 
 # Logit softcapping
-@torch.compile(fullgraph = False, dynamic = True, options = torch_compile_options)
+@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