fix test

mosaicml · Oct 22, 2023 · 834e46f · 834e46f
1 parent 1604586
commit 834e46f
Showing 1 changed file with 1 addition and 2 deletions.
diff --git a/tests/test_huggingface_flash.py b/tests/test_huggingface_flash.py
@@ -181,8 +181,7 @@ def test_flash2(model_name: str, use_flash_attention_2: bool):
         model = COMPOSER_MODEL_REGISTRY[model_cfg['name']](model_cfg, tokenizer)
 
         # check that it actually used flash attention 2
-        assert model.model.config._flash_attn_2_enabled if is_flash_v2_installed(
-        ) else not model.model.config._flash_attn_2_enabled
+        assert model.model.config._flash_attn_2_enabled if use_flash_attention_2 else not model.model.config._flash_attn_2_enabled
         attention_layer = rgetattr(
             rgetattr(model, attention_layers_attr)[0], attention_attr)
         assert isinstance(attention_layer, flash_attn_class)