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.