From 72076e53207f6243020ba263cb8f146cf442ab3a Mon Sep 17 00:00:00 2001 From: Ti-Tai Wang Date: Tue, 11 Jul 2023 12:03:12 -0700 Subject: [PATCH] Update converter registry usage in `orttraining_test_dort_custom_ops.py` (#16663) Fix Orttraining Linux Lazy Tensor CI Orttraining Linux Lazy Tensor CI is broken. The error message is AttributeError: 'OnnxRegistry' object has no attribute 'register' --- .../orttraining_test_dort_custom_ops.py | 19 +++++++++---------- 1 file changed, 9 insertions(+), 10 deletions(-) diff --git a/orttraining/orttraining/test/python/orttraining_test_dort_custom_ops.py b/orttraining/orttraining/test/python/orttraining_test_dort_custom_ops.py index d8c8d395fa..66c1bc672d 100644 --- a/orttraining/orttraining/test/python/orttraining_test_dort_custom_ops.py +++ b/orttraining/orttraining/test/python/orttraining_test_dort_custom_ops.py @@ -39,11 +39,11 @@ def custom_exporter_for_aten_add_Tensor(x, y): # exporter. # Use custom_exporter_for_aten_add_Tensor.to_function_proto() to investigate # function representing "aten::mul.Tensor". -DEFAULT_ONNX_EXPORTER_OPTIONS.onnxfunction_dispatcher.onnx_registry.register( - "aten::mul.Tensor", - DEFAULT_ONNX_EXPORTER_OPTIONS.opset_version, - custom_exporter_for_aten_add_Tensor, - True, +DEFAULT_ONNX_EXPORTER_OPTIONS.onnxfunction_dispatcher.onnx_registry.register_custom_op( + function=custom_exporter_for_aten_add_Tensor, + namespace="aten", + op_name="mul", + overload="Tensor", ) @@ -57,11 +57,10 @@ def custom_exporter_for_foo_bar_default(x): # Ask exporter to map "torch.ops.foo.bar" to # custom_exporter_for_foo_bar_default. -DEFAULT_ONNX_EXPORTER_OPTIONS.onnxfunction_dispatcher.onnx_registry.register( - "foo::bar", - DEFAULT_ONNX_EXPORTER_OPTIONS.opset_version, - custom_exporter_for_foo_bar_default, - True, +DEFAULT_ONNX_EXPORTER_OPTIONS.onnxfunction_dispatcher.onnx_registry.register_custom_op( + function=custom_exporter_for_foo_bar_default, + namespace="foo", + op_name="bar", )