ONNX export

[bpfliegel]

Hello and congrats to your work! I tried to export a trained model to ONNX, but didn't have luck, got 'TypeError: split_with_sizes(): argument 'split_sizes' (position 1) must be tuple of ints, not Tensor' for the SplAtConv2d module. The fix is easy, based on torch tracing, rchannel tensor must be cast to int when exporting at two locations in the code:

            if torch._C._get_tracing_state():
                splited = torch.split(x, int(rchannel)//self.radix, dim=1)
            else:
                splited = torch.split(x, rchannel//self.radix, dim=1)

and

            if torch._C._get_tracing_state():
                attens = torch.split(atten, int(rchannel)//self.radix, dim=1)
            else:
                attens = torch.split(atten, rchannel//self.radix, dim=1)

After that it works well. PyTorch 1.4.0 and ONNX opset version 11.

Thanks, Balint

[zhanghang1989]

Thanks for sharing!

[dengfenglai321]

torch.split

hi bpfliegel, could you tell me where to change code? thanks a lot!

[zhanghang1989]

This should be already included in the current code. Let me know if you have any difficulties in converting the model.