From bb86c9912cf762f772a4d9a6e8e0df0b167c69c1 Mon Sep 17 00:00:00 2001 From: BowenBao Date: Thu, 3 Sep 2020 13:47:22 -0700 Subject: [PATCH] fix unregister to match pytorch behavior --- 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 245d283f5a86a..57363561855e4 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)