From 3d2bcb33866694cedf558d34e8006268b678217f Mon Sep 17 00:00:00 2001 From: Justin Chu Date: Thu, 21 Jul 2022 10:47:53 -0700 Subject: [PATCH] Use unregister_custom_op_symbolic to unregister torch symbolics (#12146) Description: Use unregister_custom_op_symbolic to unregister torch symbolics Motivation and Context Fixes #11305 --- .../python/tools/pytorch_export_contrib_ops.py | 14 ++++++++------ 1 file changed, 8 insertions(+), 6 deletions(-) diff --git a/onnxruntime/python/tools/pytorch_export_contrib_ops.py b/onnxruntime/python/tools/pytorch_export_contrib_ops.py index 0483e1d56e..aaca380660 100644 --- a/onnxruntime/python/tools/pytorch_export_contrib_ops.py +++ b/onnxruntime/python/tools/pytorch_export_contrib_ops.py @@ -94,10 +94,12 @@ def register(): def unregister(): """Unregister ONNX Runtime's built-in contrib ops.""" - # TODO: replace this once PyTorch supports unregister natively. - # https://msdata.visualstudio.com/Vienna/_workitems/edit/1342343 for name in _registered_ops: - ns, kind = name.split("::") - for version in sym_help._onnx_stable_opsets: - if version >= _OPSET_VERSION and sym_registry.is_registered_op(kind, ns, version): - del sym_registry._registry[(ns, version)][kind] + try: + torch.onnx.unregister_custom_op_symbolic(name, _OPSET_VERSION) + except AttributeError: + # unregister_custom_op_symbolic is not available before PyTorch 1.12 + namespace, kind = name.split("::") + for version in sym_help._onnx_stable_opsets: + if version >= _OPSET_VERSION and sym_registry.is_registered_op(kind, namespace, version): + del sym_registry._registry[(namespace, version)][kind]