test: passing device in str type

huggingface · Nov 22, 2024 · a6251e0 · a6251e0
1 parent b4aa01a
commit a6251e0
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/tests/exporters/onnx/test_exporters_onnx_cli.py b/tests/exporters/onnx/test_exporters_onnx_cli.py
@@ -846,7 +846,7 @@ def replace_activation(model: nn.Module):
                         monolith,
                         no_post_process,
                         variant=variant,
-                        device=model.device,
+                        device=device,
                         disable_dynamic_axes_fix=not do_validation,
                         do_validation=do_validation,
                         custom_export_fn=custom_export_mock,
@@ -859,7 +859,7 @@ def replace_activation(model: nn.Module):
                 monolith,
                 no_post_process,
                 variant=variant,
-                device=model.device,
+                device=device,
                 disable_dynamic_axes_fix=True,
                 do_validation=False,
                 custom_export_fn=custom_export_mock,
@@ -873,7 +873,7 @@ def replace_activation(model: nn.Module):
                 monolith,
                 no_post_process,
                 variant=variant,
-                device=model.device,
+                device=device,
                 disable_dynamic_axes_fix=not do_validation,
                 do_validation=do_validation,
                 custom_export_fn=custom_export_mock,