Return the original forward after exporting the PyTorch model to Open…

…VINO (#522)
huggingface · Jan 18, 2024 · e00c5fc · e00c5fc
1 parent 76ce9de
commit e00c5fc
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/optimum/exporters/openvino/convert.py b/optimum/exporters/openvino/convert.py
@@ -428,6 +428,9 @@ def ts_patched_forward(*args, **kwargs):
                 compression_option=compression_option,
                 compression_ratio=compression_ratio,
             )
+        # return original forward
+        if patch_model_forward:
+            model.forward = orig_forward
         ordered_dummy_inputs = {param: dummy_inputs[param] for param in sig.parameters if param in dummy_inputs}
         ordered_input_names = list(inputs)
         flatten_inputs = flattenize_inputs(ordered_dummy_inputs.values())