fix bypass

Author: xadupre

_unittests/ut_torch_models/test_llms.py

@@ -30,8 +30,9 @@ def test_export_tiny_llm_2_bypassed(self):
         model, inputs = data["model"], data["inputs"]
         self.assertEqual({"attention_mask", "past_key_values", "input_ids"}, set(inputs))
         with bypass_export_some_errors(
-            patch_transformers=True, replace_dynamic_cache=True, verbose=10
-        ):
+            patch_transformers=True, replace_dynamic_cache=True
+        ) as modificator:
+            inputs = modificator(inputs)
             ep = torch.export.export(
                 model, (), kwargs=inputs, dynamic_shapes=data["dynamic_shapes"]
             )

_unittests/ut_torch_models/test_llms_onnx.py

@@ -18,7 +18,10 @@ def test_onnx_export_tiny_llm(self):
         data = get_tiny_llm()
         model, inputs = data["model"], data["inputs"]
         self.assertEqual({"attention_mask", "past_key_values", "input_ids"}, set(inputs))
-        with bypass_export_some_errors(patch_transformers=True, replace_dynamic_cache=True):
+        with bypass_export_some_errors(
+            patch_transformers=True, replace_dynamic_cache=True
+        ) as modificator:
+            inputs = modificator(inputs)
             ep = torch.onnx.export(
                 model,
                 (),
@@ -38,7 +41,10 @@ def test_onnx_export_tiny_llm_cdbg(self):
         data = get_tiny_llm()
         model, inputs = data["model"], data["inputs"]
         self.assertEqual({"attention_mask", "past_key_values", "input_ids"}, set(inputs))
-        with bypass_export_some_errors(patch_transformers=True, replace_dynamic_cache=True):
+        with bypass_export_some_errors(
+            patch_transformers=True, replace_dynamic_cache=True
+        ) as modificator:
+            inputs = modificator(inputs)
             onx = to_onnx(model, (), kwargs=inputs, dynamic_shapes=data["dynamic_shapes"])
         self.assert_onnx_disc(
             inspect.currentframe().f_code.co_name, onx, model, inputs, verbose=1