diff --git a/src/relax/backend/contrib/nnapi/codegen.cc b/src/relax/backend/contrib/nnapi/codegen.cc
index d9f26c417257..5e95f375d2d0 100644
--- a/src/relax/backend/contrib/nnapi/codegen.cc
+++ b/src/relax/backend/contrib/nnapi/codegen.cc
@@ -258,7 +258,9 @@ Array<runtime::Module> NNAPICompiler(Array<Function> functions, Map<String, ffi:
     auto constant_names = serializer.GetConstantNames();
     const auto pf = tvm::ffi::Function::GetGlobalRequired("runtime.nnapi_runtime_create");
     auto func_name = GetExtSymbol(func);
-    compiled_functions.push_back(pf(func_name, graph_json, constant_names));
+    auto result = pf(func_name, graph_json, constant_names);
+    tvm::runtime::Module mod = result.cast<tvm::runtime::Module>();
+    compiled_functions.push_back(mod);
   }
 
   return compiled_functions;
diff --git a/tests/python/nightly/test_nnapi/test_ops.py b/tests/python/nightly/test_nnapi/test_ops.py
index 31d584db1396..a6837d2ce5c1 100644
--- a/tests/python/nightly/test_nnapi/test_ops.py
+++ b/tests/python/nightly/test_nnapi/test_ops.py
@@ -272,7 +272,7 @@ def main(
             ) -> R.Tensor((1, 10, 1), "float32"):
                 n = T.int64()
                 with R.dataflow():
-                    t0: R.Tensor((1, 10, 15), "float32") = R.mean(i0, axis=[-1], keepdims=True)
+                    t0: R.Tensor((1, 10, 1), "float32") = R.mean(i0, axis=[-1], keepdims=True)
                     R.output(t0)
                 return t0