From a5751f75781f93aecd565e34c00ab65fc2f4f76d Mon Sep 17 00:00:00 2001 From: patrickvonplaten Date: Sun, 8 Mar 2020 22:27:00 +0100 Subject: [PATCH] fix bug with attention_mask as optional input argument --- tests/test_modeling_bart.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_modeling_bart.py b/tests/test_modeling_bart.py index 5c361d0499..6f1676ddce 100644 --- a/tests/test_modeling_bart.py +++ b/tests/test_modeling_bart.py @@ -313,7 +313,7 @@ class BartHeadTests(unittest.TestCase): config, input_ids, batch_size = self._get_config_and_data(output_past=True) attention_mask = input_ids.ne(1) lm_model = BartForConditionalGeneration(config).eval().to(torch_device).half() - lm_model.generate(input_ids, attention_mask) + lm_model.generate(input_ids, attention_mask=attention_mask) def test_prepare_bart_decoder_inputs(self): config, *_ = self._get_config_and_data(output_past=False)