Add SpaceToDepth and NMS to map_torch_types_to_onnx (#2805)

Signed-off-by: Sai Chaitanya Gajula <quic_gsaichai@quicinc.com>

Co-authored-by: Sai Chaitanya Gajula <quic_gsaichai@quicinc.com>

TrainingExtensions/torch/src/python/aimet_torch/onnx_utils.py

@@ -109,6 +109,7 @@
     nn.LogSoftmax: ['LogSoftmax'],
     nn.LSTM: ['LSTM'],
     nn.PixelShuffle: ['DepthToSpace'],
+    nn.PixelUnshuffle: ['SpaceToDepth'],
     nn.PReLU: ['PRelu'],
     nn.ReLU: ['Relu'],
     nn.ReLU6: ['Clip'],
@@ -145,7 +146,8 @@
     elementwise_ops.Tile: ['Tile'],
     elementwise_ops.TopK: ['TopK'],
     torchvision.ops.RoIPool: ['MaxRoiPool'],
-    elementwise_ops.Mean: ['ReduceMean']
+    elementwise_ops.Mean: ['ReduceMean'],
+    elementwise_ops.NonMaxSuppression: ['NonMaxSuppression'],
 }
 
 # Maps pytorch functional op string names to corresponding onnx types.