"Unsupported: ONNX export of transpose for tensor of unknown rank." with dynamic axes

[ghost]

🐛 Bug

When exporting a model with dynamic_axes, I get an error Unsupported: ONNX export of transpose for tensor of unknown rank. although the tensor itself looks correct. The model uses transpose(1,2) which is the source of the error. Please see the sample below.

To Reproduce

Steps to reproduce the behavior:

import torch

class Model(torch.nn.Module):
    def forward(self, x):
        xt = x.transpose(1,2)
        return xt

m  = Model().cuda()
i = (torch.randn(1,2,3).cuda(),)

torch.onnx.export(m, i, "model.onnx",
                  input_names=["INPUT_0"],
                  output_names=["OUTPUT_0"],
                  dynamic_axes={"INPUT_0": {0: "batch_size"},
                                "OUTPUT_0": {0: "batch_size"}})

Expected behavior

Successful model export to ONNX.

Environment

• PyTorch Version: 1.8.0a0+1606899
• OS: Ubuntu 20.04.1 LTS (x86_64)
• How you installed PyTorch: from docker nvcr.io/nvidia/pytorch:20.12-py3
• Python version: 3.8 (64-bit runtime)
• CUDA/cuDNN version: 11.1/8.0.5
• GPU models and configuration: GPU 0: TITAN V

cc @houseroad @spandantiwari @lara-hdr @BowenBao @neginraoof

[neginraoof]

Thanks for reporting this issue. Transpose shape inference logic has been recently updated in (#50163), and looks like this issue is resolved in the nightly build.
Can you confirm which build you are using? Would it be possible to test the latest nightly build?
Thanks.

[spandantiwari]

Thanks @neginraoof .

@grzegorzkarchnv - please reopen this issue if you still see this problem.

[ghost]

I confirm the issue is resolved and I can export the model successfully. Thanks for your support!

[spandantiwari]

Thanks for confirming @grzegorzkarchnv .

[Aloento]

Same issue for the ONNX 16
RuntimeError: Unsupported: ONNX export of transpose for tensor of unknown rank.
1.12.1+cu116

[N1njaG]

Same issue for the ONNX 16 RuntimeError: Unsupported: ONNX export of transpose for tensor of unknown rank. 1.12.1+cu116

Have you solved the issue？and how did you do

[fernandorovai]

I'm facing the same problem when transposing a tensor inside a for loop. Any idea on how to fix it?