diff --git a/python/tvm/micro/model_library_format.py b/python/tvm/micro/model_library_format.py index d13947527e81f..e220fa1ca5430 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)}")