Skip to content

Commit f02caa5

Browse files
Fix ShuffleNet ONNX export (#7686)
Co-authored-by: Nicolas Hug <[email protected]>
1 parent b5401b9 commit f02caa5

File tree

1 file changed

+1
-1
lines changed

1 file changed

+1
-1
lines changed

torchvision/models/shufflenetv2.py

+1-1
Original file line numberDiff line numberDiff line change
@@ -35,7 +35,7 @@ def channel_shuffle(x: Tensor, groups: int) -> Tensor:
3535
x = torch.transpose(x, 1, 2).contiguous()
3636

3737
# flatten
38-
x = x.view(batchsize, -1, height, width)
38+
x = x.view(batchsize, num_channels, height, width)
3939

4040
return x
4141

0 commit comments

Comments
 (0)