SUMIN/HuggingFace_transformer
1
0
Fork 0
Code Issues Pull Requests Actions 7 Packages Projects Releases Wiki Activity

[BART] FP16 testing fixes (#3266)

Browse Source
This commit is contained in:
Sam Shleifer
2020-03-13 19:48:26 -04:00
committed by GitHub
parent 8320feec09
commit 2bd79e23de
2 changed files with 16 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files

10
tests/test_modeling_bart.py
View File

@@ -314,10 +314,16 @@ 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
attention_mask = input_ids.ne(1).to(torch_device)
model = BartForConditionalGeneration(config).eval().to(torch_device).half()
model.generate(input_ids, attention_mask=attention_mask, do_sample=False, early_stopping=True)
@unittest.skipIf(torch_device == "cpu", "Cant do half precision")
def test_base_model_fp16(self):
config, input_ids, batch_size = self._get_config_and_data(output_past=False)
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)
lm_model(input_ids, attention_mask=attention_mask)
def test_prepare_bart_decoder_inputs(self):
config, *_ = self._get_config_and_data(output_past=False)