revert to tvm.runtime.Module

apache · Jun 16, 2022 · 06c9b51 · 06c9b51
1 parent a7560f1
commit 06c9b51
Showing 1 changed file with 2 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)}")