From d81090f05e7c3a021f0ad764d05b4e616b30748d Mon Sep 17 00:00:00 2001 From: "[[ -z $EMAIL ]] && read -e -p \"Enter your email (for git configuration): \" EMAIL" Date: Thu, 9 Jan 2025 14:43:08 -0500 Subject: [PATCH] Add sdpa to diffllama --- src/transformers/models/diffllama/modeling_diffllama.py | 1 + 1 file changed, 1 insertion(+) diff --git a/src/transformers/models/diffllama/modeling_diffllama.py b/src/transformers/models/diffllama/modeling_diffllama.py index 7d6d6af3824d..f474fe97b9be 100644 --- a/src/transformers/models/diffllama/modeling_diffllama.py +++ b/src/transformers/models/diffllama/modeling_diffllama.py @@ -595,6 +595,7 @@ class DiffLlamaPreTrainedModel(PreTrainedModel): _skip_keys_device_placement = ["past_key_values"] _supports_flash_attn_2 = True _supports_sdpa = True + _supports_flex_attn = True _supports_cache_class = True _supports_quantized_cache = True _supports_static_cache = True