From 24ea3a9ba083f2afda2efaba04bb13ced6f383ea Mon Sep 17 00:00:00 2001
From: Rishabh Thakur <quic_ristha@quicinc.com>
Date: Tue, 19 Dec 2023 12:32:02 +0530
Subject: [PATCH] Add torch to ONNX mapping for InstanceNorm (#2609)

Signed-off-by: Sai Chaitanya Gajula <quic_gsaichai@quicinc.com>
---
 TrainingExtensions/torch/src/python/aimet_torch/onnx_utils.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/TrainingExtensions/torch/src/python/aimet_torch/onnx_utils.py b/TrainingExtensions/torch/src/python/aimet_torch/onnx_utils.py
index 318c7782fd0..3c9a8e4ef60 100644
--- 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'],