Adding mapping for mean in the torch to onnx op naming dict in onnx_u…

…tils. (#2624) Signed-off-by: Sayanta Mukherjee <[email protected]>
quic · Dec 28, 2023 · 48e8781 · 48e8781
1 parent ab0ed37
commit 48e8781
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 2 deletions.
diff --git a/TrainingExtensions/tensorflow/src/python/aimet_tensorflow/keras/model_preparer.py b/TrainingExtensions/tensorflow/src/python/aimet_tensorflow/keras/model_preparer.py
@@ -50,7 +50,7 @@
     from keras.engine.functional import Functional # pylint: disable=import-error
     from keras.engine.keras_tensor import KerasTensor # pylint: disable=import-error
     from keras.layers.core.tf_op_layer import TFOpLambda # pylint: disable=import-error
-    from keras.layers.merging.base_merge import _Merge as MergeLayersParentClass # pylint: disable=ungrouped-imports
+    from keras.layers.merging.base_merge import _Merge as MergeLayersParentClass # pylint: disable=ungrouped-imports, import-error
 else:
     # Ignore pylint errors due to conditional imports
     from tensorflow.python.keras.engine.base_layer_utils import is_subclassed # pylint: disable=ungrouped-imports

diff --git a/TrainingExtensions/torch/src/python/aimet_torch/onnx_utils.py b/TrainingExtensions/torch/src/python/aimet_torch/onnx_utils.py
@@ -144,7 +144,8 @@
     elementwise_ops.Subtract: ['Sub'],
     elementwise_ops.Tile: ['Tile'],
     elementwise_ops.TopK: ['TopK'],
-    torchvision.ops.RoIPool: ['MaxRoiPool']
+    torchvision.ops.RoIPool: ['MaxRoiPool'],
+    elementwise_ops.Mean: ['ReduceMean']
 }
 
 # Maps pytorch functional op string names to corresponding onnx types.