fix test

huggingface · Sep 18, 2023 · 8d102f7 · 8d102f7
1 parent dfe7e5e
commit 8d102f7
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 1 deletion.
diff --git a/optimum/onnxruntime/modeling_decoder.py b/optimum/onnxruntime/modeling_decoder.py
@@ -1036,6 +1036,7 @@ def _from_transformers(
             task=task,
             do_validation=False,
             no_post_process=False,
+            legacy=False,
             subfolder=subfolder,
             revision=revision,
             cache_dir=cache_dir,

diff --git a/tests/onnxruntime/test_modeling.py b/tests/onnxruntime/test_modeling.py
@@ -2049,7 +2049,7 @@ def test_merge_from_onnx_and_save(self, model_arch):
             self.skipTest("Unsupported export case")
 
         with tempfile.TemporaryDirectory() as tmpdir:
-            main_export(model_id, tmpdir, task=task)
+            main_export(model_id, tmpdir, task=task, legacy=True)
 
             model = ORTModelForCausalLM.from_pretrained(tmpdir)