The register_custom_op_symbolic function, part of torch.onnx does not see to be working properly for me. I am following this simple example found in the documentation using a contributed operator thats part of onnxruntime-gpu. I noted that the ONNX unit tests do not check the functionality using a custom namespace, so I am assuming this doc example is correct. I tried to do some digging into the torch/_ops.py but was not able to follow how ops are dynamically created.

I am running in the Nvidia PyTorch container 23.06, only additional install I do is pip install onnxruntime-gpu.

I cannot get the symbolic operation to register properly (I can't call it from torch.ops). Here's a minimal example (based on the example in the docs linked above):

The error is related to accessing the symbolic ops in PyTorch, so I am guessing its an issue with this register function. I would expect torch.ops.custom_ops.foo_forward to be callable since I registered the operator.

Any help would be appreciated, thanks!