From 7f65daa2e155ecdd8594e19862dac8b322ed3b73 Mon Sep 17 00:00:00 2001 From: Patrick von Platen Date: Tue, 4 Aug 2020 13:02:25 +0200 Subject: [PATCH] fix reformer fp16 (#6237) --- tests/test_modeling_reformer.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_modeling_reformer.py b/tests/test_modeling_reformer.py index a56b99c143..e5d07d8eb6 100644 --- a/tests/test_modeling_reformer.py +++ b/tests/test_modeling_reformer.py @@ -389,7 +389,7 @@ class ReformerModelTester: model.to(torch_device) model.half() model.eval() - output = model(input_ids, attention_mask=input_mask)["last_input_state"] + output = model(input_ids, attention_mask=input_mask)["last_hidden_state"] self.parent.assertFalse(torch.isnan(output).any().item()) def create_and_check_reformer_model_generate(self, config, input_ids, input_mask, choice_labels):