Fix fp16 ONNX export test (#1373)

fix test
huggingface · Sep 8, 2023 · 058c72d · 058c72d
1 parent b7a11b9
commit 058c72d
Showing 1 changed file with 9 additions and 2 deletions.
diff --git a/tests/exporters/onnx/test_exporters_onnx_cli.py b/tests/exporters/onnx/test_exporters_onnx_cli.py
@@ -419,7 +419,14 @@ def test_stable_diffusion(self):
     @slow
     @pytest.mark.run_slow
     def test_export_on_fp16(
-        self, test_name: str, model_type: str, model_name: str, task: str, monolith: bool, no_post_process: bool
+        self,
+        test_name: str,
+        model_type: str,
+        model_name: str,
+        task: str,
+        variant: str,
+        monolith: bool,
+        no_post_process: bool,
     ):
         # TODO: refer to https://github.com/pytorch/pytorch/issues/95377
         if model_type == "yolos":
@@ -446,7 +453,7 @@ def test_export_on_fp16(
         if model_type == "ibert":
             self.skipTest("ibert can not be supported in fp16")
 
-        self._onnx_export(model_name, task, monolith, no_post_process, fp16=True)
+        self._onnx_export(model_name, task, monolith, no_post_process, variant=variant, fp16=True, device="cuda")
 
     @parameterized.expand(
         [