Fix ShuffleNet ONNX export (#7686)

Co-authored-by: Nicolas Hug <nh.nicolas.hug@gmail.com>
pytorch · Jun 22, 2023 · f02caa5 · f02caa5
1 parent b5401b9
commit f02caa5
Showing 1 changed file with 1 addition and 1 deletion.
diff --git 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