[ONNX] Fix ShuffleNetV2 model export issue.

test/test_onnx.py

@@ -478,6 +478,17 @@ def test_keypoint_rcnn(self):
                        dynamic_axes={"images_tensors": [0, 1, 2]},
                        tolerate_small_mismatch=True)
 
+    def test_shufflenet_v2_dynamic_axes(self):
+        model = models.shufflenet_v2_x0_5(pretrained=True)
+        dummy_input = torch.randn(1, 3, 224, 224, requires_grad=True)
+        test_inputs = torch.cat([dummy_input, dummy_input, dummy_input], 0)
+
+        self.run_model(model, [(dummy_input,), (test_inputs,)],
+                       input_names=["input_images"],
+                       output_names=["output"],
+                       dynamic_axes={"input_images": {0: 'batch_size'}, "output": {0: 'batch_size'}},
+                       tolerate_small_mismatch=True)
+
 
 if __name__ == '__main__':
     unittest.main()

torchvision/models/shufflenetv2.py

@@ -19,7 +19,7 @@
 
 
 def channel_shuffle(x: Tensor, groups: int) -> Tensor:
-    batchsize, num_channels, height, width = x.data.size()
+    batchsize, num_channels, height, width = x.shape[0], x.shape[1], x.shape[2], x.shape[3]
     channels_per_group = num_channels // groups
 
     # reshape