fix bug in export_to_onnx API (#1166)

intel · Aug 18, 2022 · 158c7f4 · 158c7f4
1 parent a28705c
commit 158c7f4
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/neural_compressor/model/torch_model.py b/neural_compressor/model/torch_model.py
@@ -324,7 +324,7 @@ def graph_info(self):
 
     def export_to_jit(self, example_inputs=None):
         if example_inputs is not None:
-            if isinstance(input, dict) or isinstance(input, UserDict):
+            if isinstance(example_inputs, dict) or isinstance(example_inputs, UserDict):
                 example_inputs = tuple(example_inputs.values())
         else:
             logger.warning("Please provide example_inputs for jit.trace")
@@ -357,8 +357,8 @@ def export_to_fp32_onnx(
         fp32_model=None,
     ):
         example_input_names = ['input']
-        if isinstance(input, dict) or isinstance(input, UserDict):
-            example_input_names = list(input.keys())
+        if isinstance(example_inputs, dict) or isinstance(example_inputs, UserDict):
+            example_input_names = list(example_inputs.keys())
         model = self.model
         if fp32_model:
             model = fp32_model