From 73456f10cd70554db7cd926dff9fd0fc83555fc9 Mon Sep 17 00:00:00 2001 From: Bowen Bao Date: Thu, 3 Sep 2020 16:32:42 -0700 Subject: [PATCH] Fix contrib ops unregister to match pytorch behavior (#5052) --- tools/python/register_custom_ops_pytorch_exporter.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/tools/python/register_custom_ops_pytorch_exporter.py b/tools/python/register_custom_ops_pytorch_exporter.py index 245d283f5a..5736356185 100644 --- a/tools/python/register_custom_ops_pytorch_exporter.py +++ b/tools/python/register_custom_ops_pytorch_exporter.py @@ -43,10 +43,14 @@ def unregister_custom_op(): import torch.onnx.symbolic_registry as sym_registry + # TODO: replace this once PyTorch supports unregister natively. def unregister(name, opset_version): ns, kind = name.split("::") - if sym_registry.is_registered_op(kind, ns, opset_version): - del sym_registry._registry[(ns, opset_version)][kind] + from torch.onnx.symbolic_helper import _onnx_stable_opsets + + for version in _onnx_stable_opsets: + if version >= opset_version and sym_registry.is_registered_op(kind, ns, version): + del sym_registry._registry[(ns, version)][kind] unregister('::inverse', _onnx_opset_version) unregister('::gelu', _onnx_opset_version)