diff --git a/python/akg/composite/build_module.py b/python/akg/composite/build_module.py
index dfa9428cbb8cf454c0c3ad1fddee4da20181c349..3f725fcf99f98bcc0c83fb9f4587da59b2c80e35 100644
--- a/python/akg/composite/build_module.py
+++ b/python/akg/composite/build_module.py
@@ -15,7 +15,9 @@
 # limitations under the License.
 
 """build module"""
+import os
 import json
+import akg
 from akg import tvm
 from akg.tvm import _api_internal
 from .repository import __all__ as repository
@@ -178,6 +180,6 @@ def build_cuda(outputs, args, sch_name, kernel_name):
         s = scheduler[sch_name](outputs)
         dump_ir = os.getenv('MS_AKG_DUMP_IR') == "on"
         with tvm.build_config(dump_pass_ir = dump_ir):
-            mod = tvm.build(s, args, cuda, name = kernel_name)
+            mod = akg.build(s, list(args), "cuda", name = kernel_name)
             dump_cuda_meta.dump(mod, kernel_name, s, list(args))
             return mod
diff --git a/python/akg/ms/message.py b/python/akg/ms/message.py
index 94e533a0d3a4fd30649aa73ec3f053656615a527..3a856a4b630119505940c3868c38046c34363df4 100644
--- a/python/akg/ms/message.py
+++ b/python/akg/ms/message.py
@@ -20,7 +20,7 @@ import json
 import json.decoder as jd
 import logging
 import traceback
-import os.path
+import os
 from pathlib import Path
 
 import akg.tvm