riaz.somc/transformers
1
0
Fork 0
mirror of https://github.com/huggingface/transformers.git synced 2025-07-31 02:02:21 +06:00
Code Issues Actions 33 Packages Projects Releases Wiki Activity

make all tensors half precision

Browse Source
This commit is contained in:
Patrick von Platen 2020-03-11 12:15:38 +01:00
parent a332cc9f7f
commit bc9d5d917c
1 changed files with 2 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
tests/test_modeling_bart.py
View File

@ -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)