From ac303eae46874e48854b60626b9db7a407688e3d Mon Sep 17 00:00:00 2001 From: Patrick von Platen Date: Wed, 11 Mar 2020 12:24:30 +0100 Subject: [PATCH] fix problem with half --- tests/test_modeling_bart.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/test_modeling_bart.py b/tests/test_modeling_bart.py index 2bcc2c4b45..b885ccf1b6 100644 --- a/tests/test_modeling_bart.py +++ b/tests/test_modeling_bart.py @@ -314,8 +314,8 @@ class BartHeadTests(unittest.TestCase): @unittest.skipIf(torch_device == "cpu", "Cant do half precision") def test_generate_fp16(self): config, input_ids, batch_size = self._get_config_and_data(output_past=True) - input_ids = input_ids.half() - attention_mask = input_ids.ne(1).to(torch_device).half() + input_ids = input_ids + attention_mask = input_ids.ne(1).to(torch_device) lm_model = BartForConditionalGeneration(config).eval().to(torch_device).half() lm_model.generate(input_ids, attention_mask=attention_mask)