SUMIN/HuggingFace_transformer
1
0
Fork 0
Code Issues Pull Requests Actions 7 Packages Projects Releases Wiki Activity

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

Browse Source 
explicit use_cache=True
This commit is contained in:
fxmarty
2023-11-21 17:53:10 +01:00
committed by GitHub
parent c770600fde
commit 7f04373865
2 changed files with 10 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
tests/models/mistral/test_modeling_mistral.py
View File

@@ -436,7 +436,11 @@ class MistralModelTest(ModelTesterMixin, GenerationTesterMixin, PipelineTesterMi
# Just test that a large cache works as expected # Just test that a large cache works as expected
_ = model.generate( _ = 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 @require_flash_attn

6
tests/test_modeling_common.py
View File

@@ -3166,7 +3166,11 @@ class ModelTesterMixin:
# Just test that a large cache works as expected # Just test that a large cache works as expected
_ = model.generate( _ = 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 @require_flash_attn