Skip to content

GitLab

"test/git@developer.sourcefind.cn:OpenDAS/torchaudio.git" did not exist on "df48ba36826b838ec8d986343ad0ac42964314d0"
Commit bc9d5d91 authored by Patrick von Platen's avatar Patrick von Platen
Browse files

make all tensors half precision

parent a332cc9f
Hide whitespace changes
Inline Side-by-side
Showing with 2 additions and 1 deletion
tests/test_modeling_bart.py
View file @ bc9d5d91
...@@ -314,7 +314,8 @@ class BartHeadTests(unittest.TestCase): ...@@ -314,7 +314,8 @@ class BartHeadTests(unittest.TestCase):
@unittest.skipIf(torch_device == "cpu", "Cant do half precision") @unittest.skipIf(torch_device == "cpu", "Cant do half precision")
def test_generate_fp16(self): def test_generate_fp16(self):
config, input_ids, batch_size = self._get_config_and_data(output_past=True) 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 = BartForConditionalGeneration(config).eval().to(torch_device).half()
lm_model.generate(input_ids, attention_mask=attention_mask) lm_model.generate(input_ids, attention_mask=attention_mask)
......
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