From d533c7e9b9a5122cb1381702b99a3aff0aa2e0c4 Mon Sep 17 00:00:00 2001 From: Funtowicz Morgan Date: Wed, 15 Jul 2020 16:11:22 +0200 Subject: [PATCH] [fix] T5 ONNX test: model.to(torch_device) (#5769) Signed-off-by: Morgan Funtowicz --- tests/test_modeling_t5.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_modeling_t5.py b/tests/test_modeling_t5.py index bcf5e76481..6c4054dc46 100644 --- 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,