fix

huggingface · Feb 4, 2025 · f1d2e1a · f1d2e1a
1 parent 4d96f42
commit f1d2e1a
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/tests/test_modeling_common.py b/tests/test_modeling_common.py
@@ -4459,7 +4459,8 @@ def test_flash_attn_2_can_dispatch_composite_models(self):
                     for name, submodule in model_fa2.named_modules():
                         class_name = submodule.__class__.__name__
                         if (
-                            class_name.endswith("Attention")
+                            "Attention" in class_name
+                            and getattr(submodule, "config", None)
                             and submodule.config._attn_implementation == "flash_attention_2"
                         ):
                             has_fa2 = True