minor fix

Author: yarden-sony

model_compression_toolkit/exporter/model_exporter/pytorch/fakely_quant_onnx_pytorch_exporter.py

@@ -100,15 +100,16 @@ def export(self, output_names=None) -> None:
             model_output = self.model(*model_input) if isinstance(model_input, (list, tuple)) else self.model(
                 model_input)
 
+            input_names = [n.name for n in self.model.node_sort if n.type == DummyPlaceHolder]
+            dynamic_axes = {name: {0: 'batch_size'} for name in input_names}
             if output_names is None:
                 # Determine number of outputs and prepare output_names and dynamic_axes
                 if isinstance(model_output, (list, tuple)):
                     output_names = [f"output_{i}" for i in range(len(model_output))]
-                    dynamic_axes = {'input': {0: 'batch_size'}}
                     dynamic_axes.update({name: {0: 'batch_size'} for name in output_names})
                 else:
                     output_names = ['output']
-                    dynamic_axes = {'input': {0: 'batch_size'}, 'output': {0: 'batch_size'}}
+                    dynamic_axes.update({'output': {0: 'batch_size'}})
             else:
                 if isinstance(model_output, (list, tuple)):
                     num_of_outputs = len(model_output)
@@ -117,12 +118,8 @@ def export(self, output_names=None) -> None:
                 assert len(output_names) == num_of_outputs, (f"Mismatch between number of requested output names "
                                                              f"({output_names}) and model output count "
                                                              f"({num_of_outputs}):\n")
-                dynamic_axes = {'input': {0: 'batch_size'}}
                 dynamic_axes.update({name: {0: 'batch_size'} for name in output_names})
-
-            input_names = [n.name for n in self.model.node_sort if n.type == DummyPlaceHolder]
-            dynamic_axes.update({name: {0: 'batch_size'} for name in input_names})
-
+            dynamic_axes.update({"input": {0: 'batch_size'}})
             if hasattr(self.model, 'metadata'):
                 onnx_bytes = BytesIO()
                 torch.onnx.export(self.model,