Generate: pin number of beams in BART test (#22763)

9af845af · Joao Gante · GitHub · 66b15efb · 9af845af
Unverified Commit 9af845af authored Apr 14, 2023 by Joao Gante Committed by GitHub Apr 14, 2023
Hide whitespace changes
Inline Side-by-side

Showing with 1 addition and 1 deletion

tests/models/bart/test_modeling_bart.py tests/models/bart/test_modeling_bart.py +1 -1

No files found.
--- a/tests/models/bart/test_modeling_bart.py
+++ b/tests/models/bart/test_modeling_bart.py
@@ -1230,7 +1230,7 @@ class BartModelIntegrationTests(unittest.TestCase):
            article, add_special_tokens=False, truncation=True, max_length=512, return_tensors="pt"
        ).input_ids.to(torch_device)

-        outputs = bart_model.generate(input_ids, penalty_alpha=0.5, top_k=5, max_length=64)
+        outputs = bart_model.generate(input_ids, penalty_alpha=0.5, top_k=5, max_length=64, num_beams=1)
        generated_text = bart_tokenizer.batch_decode(outputs, skip_special_tokens=True)

        self.assertListEqual(