Beef up Llama tests (#22314)

* tmp commit

* beef up llama tests
This commit is contained in:
Joao Gante
2023-03-22 15:20:48 +00:00
committed by GitHub
parent 12febc20db
commit fd3eb3e3cd
2 changed files with 15 additions and 18 deletions

View File

@@ -1463,10 +1463,10 @@ class GenerationTesterMixin:
attention_names = ["encoder_attentions", "decoder_attentions", "cross_attentions"]
for model_class in self.all_generative_model_classes:
config, input_ids, attention_mask, max_length = self._get_input_ids_and_config()
model = model_class(config).to(torch_device)
# We want to test only encoder-decoder models
if not config.is_encoder_decoder:
continue
model = model_class(config).to(torch_device)
head_masking = {
"head_mask": torch.zeros(config.encoder_layers, config.encoder_attention_heads, device=torch_device),