From 58fc8f97a39fc9c9e71e141763a445549828f684 Mon Sep 17 00:00:00 2001 From: patrickvonplaten Date: Fri, 6 Mar 2020 00:35:47 +0100 Subject: [PATCH] fix renaming problem --- 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 565335e69b..374c7b840f 100644 --- a/tests/test_modeling_bart.py +++ b/tests/test_modeling_bart.py @@ -300,7 +300,7 @@ class BartHeadTests(unittest.TestCase): def test_generate_fp16(self): config, input_ids, batch_size = self._get_config_and_data(output_past=True) attention_mask = input_ids.ne(1) - lm_model = BartForMaskedLM(config).eval().to(torch_device).half() + lm_model = BartForConditionalGeneration(config).eval().to(torch_device).half() lm_model.generate(input_ids, attention_mask)