Explicitely specify `use_cache=True` in Flash Attention tests (#27635)

explicit use_cache=True

Explicitely specify `use_cache=True` in Flash Attention tests (#27635)
explicit use_cache=True
7f043738 · fxmarty · GitHub · c770600f · 7f043738 · 7f043738
Unverified Commit 7f043738 authored Nov 21, 2023 by fxmarty Committed by GitHub Nov 22, 2023
Show whitespace changes
Inline Side-by-side

Showing with 10 additions and 2 deletions

tests/models/mistral/test_modeling_mistral.py tests/models/mistral/test_modeling_mistral.py +5 -1

tests/test_modeling_common.py tests/test_modeling_common.py +5 -1

No files found.
--- a/tests/models/mistral/test_modeling_mistral.py
+++ b/tests/models/mistral/test_modeling_mistral.py
@@ -436,7 +436,11 @@ class MistralModelTest(ModelTesterMixin, GenerationTesterMixin, PipelineTesterMi

                # Just test that a large cache works as expected
                _ = model.generate(
-                    dummy_input, attention_mask=dummy_attention_mask, max_new_tokens=max_new_tokens, do_sample=False
+                    dummy_input,
+                    attention_mask=dummy_attention_mask,
+                    max_new_tokens=max_new_tokens,
+                    do_sample=False,
+                    use_cache=True,
                )

    @require_flash_attn

--- a/tests/test_modeling_common.py
+++ b/tests/test_modeling_common.py
@@ -3166,7 +3166,11 @@ class ModelTesterMixin:

                # Just test that a large cache works as expected
                _ = model.generate(
-                    dummy_input, attention_mask=dummy_attention_mask, max_new_tokens=max_new_tokens, do_sample=False
+                    dummy_input,
+                    attention_mask=dummy_attention_mask,
+                    max_new_tokens=max_new_tokens,
+                    do_sample=False,
+                    use_cache=True,
                )

    @require_flash_attn