[fix] T5 ONNX test: model.to(torch_device) (#5769)

Signed-off-by: Morgan Funtowicz <morgan@huggingface.co>

[fix] T5 ONNX test: model.to(torch_device) (#5769)
Signed-off-by: Morgan Funtowicz <morgan@huggingface.co>
d533c7e9 · Funtowicz Morgan · GitHub · d0486c8b · d533c7e9
Unverified Commit d533c7e9 authored Jul 15, 2020 by Funtowicz Morgan Committed by GitHub Jul 15, 2020
Show whitespace changes
Inline Side-by-side

Showing with 1 addition and 1 deletion

tests/test_modeling_t5.py tests/test_modeling_t5.py +1 -1

No files found.
--- a/tests/test_modeling_t5.py
+++ b/tests/test_modeling_t5.py
@@ -336,7 +336,7 @@ class T5ModelTest(ModelTesterMixin, unittest.TestCase):

        config_and_inputs = self.model_tester.prepare_config_and_inputs()
        config_and_inputs[0].return_tuple = True
-        model = T5Model(config_and_inputs[0])
+        model = T5Model(config_and_inputs[0]).to(torch_device)
        with tempfile.TemporaryDirectory() as tmpdirname:
            torch.onnx.export(
                model, config_and_inputs[1], f"{tmpdirname}/t5_test.onnx", export_params=True, opset_version=9,