revert to tvm.runtime.Module

apache · Jun 16, 2022 · 36e38f9 · 36e38f9
1 parent a7560f1
commit 36e38f9
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 2 deletions.
diff --git 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
@@ -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}