I have a SoftmaxTree implementation in C++ extension, that I cannot convert to ONNX. I can change the implementation to use recursive jit script, but before that I wanted to see if then it can be exported to ONNX.
I have tried simple @torch.jit.script with an if-statement can be exported and results in an if node in hte graph. But anything more complicated I am not sure.
It depends on your code, the TorchScript compiler always makes the same type of graph, and there are various limitations on what can be exported to ONNX. Looking at the link you posted, TorchScript currently doesn’t support custom autograd Functions, but this is something we are working on supporting.