fix patching order

huggingface · Jul 11, 2024 · d0e188f · d0e188f
1 parent f742ee6
commit d0e188f
Showing 1 changed file with 3 additions and 5 deletions.
diff --git a/optimum/exporters/openvino/convert.py b/optimum/exporters/openvino/convert.py
@@ -347,11 +347,6 @@ def export_pytorch(
             patcher = config.patch_model_for_export(model, model_kwargs=model_kwargs)
             patched_forward = patcher.patched_forward
 
-            if patch_16bit_model:
-                from openvino.frontend.pytorch.patch_model import __make_16bit_traceable
-
-                __make_16bit_traceable(model)
-
             @functools.wraps(patched_forward)
             def ts_patched_forward(*args, **kwargs):
                 for i in range(len(dict_inputs)):
@@ -365,6 +360,9 @@ def ts_patched_forward(*args, **kwargs):
             patcher.patched_forward = ts_patched_forward
 
             with patcher:
+                if patch_16bit_model:
+                    from openvino.frontend.pytorch.patch_model import __make_16bit_traceable
+                    __make_16bit_traceable(model)
                 check_dummy_inputs_are_allowed(model, dummy_inputs)
                 sig = inspect.signature(model.forward) if hasattr(model, "forward") else inspect.signature(model.call)
                 inputs = config.ordered_inputs(model)