From 3056c8f9cfad105d5a129150cc710648ce36ad85 Mon Sep 17 00:00:00 2001 From: Duc Nguyen Date: Wed, 21 Jun 2023 15:35:42 +0700 Subject: [PATCH] used explicit shape rather than -1 to avoid onnx export failure --- 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