Fix aot failed tests

apache · Sep 10, 2021 · b965443 · b965443
1 parent 5b3fb11
commit b965443
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/src/relay/backend/aot_executor_codegen.cc b/src/relay/backend/aot_executor_codegen.cc
@@ -666,8 +666,11 @@ class AOTExecutorCodegen : public MixedModeVisitor {
         Downcast<tir::PrimFunc>(mod_run->Lookup(::tvm::runtime::symbol::tvm_run_func_suffix)),
         workspace_byte_alignment);
 
+    lowered_mod = WithAttr(lowered_mod, "main_func_info", func_info);
+
     Optional<backend::FunctionInfo> main_func_info =
         lowered_mod->GetAttr<backend::FunctionInfo>("main_func_info");
+
     ICHECK(main_func_info) << "The attribute \"main_func_info\" should be set at this point.";
     main_func_info.value()->workspace_sizes.Set(target_host_, main_workspace_size);
     function_metadata_.Set(runtime::symbol::tvm_module_main, main_func_info.value());