Skip to content

GitLab

Unverified Commit 9af845af authored by Joao Gante's avatar Joao Gante Committed by GitHub
Browse files

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

parent 66b15efb
Show whitespace changes
Inline Side-by-side
Showing with 1 addition and 1 deletion
tests/models/bart/test_modeling_bart.py
View file @ 9af845af
...@@ -1230,7 +1230,7 @@ class BartModelIntegrationTests(unittest.TestCase): ...@@ -1230,7 +1230,7 @@ class BartModelIntegrationTests(unittest.TestCase):
article, add_special_tokens=False, truncation=True, max_length=512, return_tensors="pt" article, add_special_tokens=False, truncation=True, max_length=512, return_tensors="pt"
).input_ids.to(torch_device) ).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) generated_text = bart_tokenizer.batch_decode(outputs, skip_special_tokens=True)
self.assertListEqual( self.assertListEqual(
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment