From ce621b7af32b44ee2c6bbce5f9d302b6e0bc677e Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Wed, 6 Nov 2024 15:56:26 -0800 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 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