From 36e38f9db7181e5736c7ede7ac88be8c23e0301d Mon Sep 17 00:00:00 2001 From: Mehrdad Hessar Date: Tue, 14 Jun 2022 14:56:56 -0700 Subject: [PATCH] revert to tvm.runtime.Module --- python/tvm/micro/model_library_format.py | 4 ++-- tests/python/relay/aot/test_crt_aot.py | 1 + 2 files changed, 3 insertions(+), 2 deletions(-) diff --git a/python/tvm/micro/model_library_format.py b/python/tvm/micro/model_library_format.py index d13947527e81..e220fa1ca543 100644 --- a/python/tvm/micro/model_library_format.py +++ b/python/tvm/micro/model_library_format.py @@ -71,7 +71,7 @@ def generate_c_interface_header( def _populate_codegen_dir( mods: typing.Union[ typing.List[executor_factory.ExecutorFactoryModule], - typing.List[build_module.OperatorModule], + typing.List[tvm.runtime.Module], ], codegen_dir: str, ): @@ -92,7 +92,7 @@ def _populate_codegen_dir( for mod in mods: if isinstance(mod, executor_factory.ExecutorFactoryModule): lib = mod.lib - elif isinstance(mod, build_module.OperatorModule): + elif isinstance(mod, tvm.runtime.Module): lib = mod else: raise RuntimeError(f"Not supported module type: {type(mod)}") diff --git a/tests/python/relay/aot/test_crt_aot.py b/tests/python/relay/aot/test_crt_aot.py index a46e7925fe6f..987d425aa63d 100644 --- a/tests/python/relay/aot/test_crt_aot.py +++ b/tests/python/relay/aot/test_crt_aot.py @@ -76,6 +76,7 @@ def test_error_c_interface_with_packed_api(): @parametrize_aot_options def test_conv_with_params(interface_api, use_unpacked_api, test_runner): + """Tests compilation of convolution with parameters""" mod = get_conv2d_relay_module() main_func = mod["main"] shape_dict = {p.name_hint: p.checked_type.concrete_shape for p in main_func.params}