From 7ea63955c99c36d0559002970eeb967ad9a71033 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Tue, 3 Sep 2024 13:51:31 -0700 Subject: [PATCH] Update flex_attention.py --- unsloth/kernels/flex_attention.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/unsloth/kernels/flex_attention.py b/unsloth/kernels/flex_attention.py index 9a2054c5..9cf999e2 100644 --- 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