[Pytorch] fix translation of transpose when axis argument is as a list (#5451)

python/tvm/relay/frontend/pytorch.py

@@ -923,7 +923,7 @@ def _transpose(prelude):
             axes[src] = dst
             axes[dst] = src
         else:
-            axes = inputs[1]
+            axes = _infer_shape(inputs[1], prelude.mod)
         return _op.transform.transpose(data, axes)
     return _impl
 

tests/python/frontend/pytorch/test_forward.py

@@ -767,9 +767,14 @@ def test_forward_transpose():
         def forward(self, *args):
             return args[0].transpose(-2, -1)
 
+    class Transpose3(Module):
+        def forward(self, *args):
+            return args[0].permute(0,2,3,1)
+
     input_data = torch.rand(input_shape).float()
     verify_model(Transpose1().float().eval(), input_data=input_data)
     verify_model(Transpose2().float().eval(), input_data=input_data)
+    verify_model(Transpose3().float().eval(), input_data=input_data)
 
 def test_forward_size():
     torch.set_grad_enabled(False)