Add torch to ONNX mapping for InstanceNorm (#2609)

Signed-off-by: Sai Chaitanya Gajula <quic_gsaichai@quicinc.com>
quic · Dec 19, 2023 · 24ea3a9 · 24ea3a9
1 parent 486ec81
commit 24ea3a9
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/TrainingExtensions/torch/src/python/aimet_torch/onnx_utils.py b/TrainingExtensions/torch/src/python/aimet_torch/onnx_utils.py
@@ -102,6 +102,8 @@
     nn.MaxPool2d: ['MaxPool'],
     nn.MaxPool3d: ['MaxPool'],
     nn.LayerNorm: ['LayerNorm'],        # Not a supported op in ONNX, adding this entry for usage by Connected Graph
+    nn.InstanceNorm2d: ['InstanceNormalization'],
+    nn.InstanceNorm1d: ['InstanceNormalization'],
     nn.LeakyReLU: ['LeakyRelu'],
     nn.Linear: ['Gemm', 'MatMul'],
     nn.LogSoftmax: ['LogSoftmax'],