From bc9d5d917c18689aeb37c67d030bac7ccda2cfc8 Mon Sep 17 00:00:00 2001 From: Patrick von Platen Date: Wed, 11 Mar 2020 12:15:38 +0100 Subject: [PATCH] make all tensors half precision --- tests/test_modeling_bart.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/tests/test_modeling_bart.py b/tests/test_modeling_bart.py index ada32ab647..2bcc2c4b45 100644 --- a/tests/test_modeling_bart.py +++ b/tests/test_modeling_bart.py @@ -314,7 +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) - attention_mask = input_ids.ne(1) + input_ids = input_ids.half() + attention_mask = input_ids.ne(1).to(torch_device).half() lm_model = BartForConditionalGeneration(config).eval().to(torch_device).half() lm_model.generate(input_ids, attention_mask=attention_mask)