style

Author: Cyrilvallez

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 = False
     _supports_cache_class = True
     _supports_quantized_cache = True
     _supports_static_cache = True

src/transformers/models/diffllama/modular_diffllama.py

@@ -431,8 +431,7 @@ def __init__(self, config: DiffLlamaConfig, layer_idx: int):
 
 
 class DiffLlamaPreTrainedModel(LlamaPreTrainedModel):
-    pass
-
+    _supports_flex_attn = False
 
 class DiffLlamaModel(LlamaModel):
     pass