saymrwulf/onnxruntime
1
0
Fork 0
mirror of https://github.com/saymrwulf/onnxruntime.git synced 2026-05-19 21:32:23 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Fix contrib ops unregister to match pytorch behavior (#5052)

Browse source
This commit is contained in:
Bowen Bao 2020-09-03 16:32:42 -07:00 committed by GitHub
parent d7502eff8f
commit 73456f10cd
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
1 changed files with 6 additions and 2 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

8
tools/python/register_custom_ops_pytorch_exporter.py
View file

@ -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)