I have already created this issue.
Many times we want to know if we are in the middle of exporting to ONNX to select a different path (say MemoryEfficientSwish vs Swish). Passing a boolean around all the time is not a great idea. So far I have been using torch._C._get_tracing_state()
but it is private API.
I just noticed there is torch.onnx.is_in_onnx_export()
. Is it equivalent to torch._C._get_tracing_state()
?