From f02caa53307c4c157c210c9fbe3fffb97ac2e635 Mon Sep 17 00:00:00 2001
From: Duc Nguyen <justanhduc@users.noreply.github.com>
Date: Thu, 22 Jun 2023 18:01:51 +0700
Subject: [PATCH] Fix ShuffleNet ONNX export (#7686)

Co-authored-by: Nicolas Hug <nh.nicolas.hug@gmail.com>
---
 torchvision/models/shufflenetv2.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/torchvision/models/shufflenetv2.py b/torchvision/models/shufflenetv2.py
index 52b85244b3d..3f3322b7a88 100644
--- a/torchvision/models/shufflenetv2.py
+++ b/torchvision/models/shufflenetv2.py
@@ -35,7 +35,7 @@ def channel_shuffle(x: Tensor, groups: int) -> Tensor:
     x = torch.transpose(x, 1, 2).contiguous()
 
     # flatten
-    x = x.view(batchsize, -1, height, width)
+    x = x.view(batchsize, num_channels, height, width)
 
     return x