fix crt test

apache · Sep 14, 2021 · 06d7331 · 06d7331
1 parent f95c459
commit 06d7331
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 4 deletions.
diff --git a/python/tvm/micro/build.py b/python/tvm/micro/build.py
@@ -23,6 +23,7 @@
 import pathlib
 import contextlib
 
+from typing import Union
 from .._ffi import libinfo
 from .. import rpc as _rpc
 
@@ -68,17 +69,19 @@ class AutoTvmModuleLoader:
 
     Parameters
     ----------
-    template_project_dir : str
+    template_project_dir : Union[pathlib.Path, str]
         project template path
 
     project_options : dict
         project generation option
     """
 
-    def __init__(self, template_project_dir: str, project_options: dict = None):
+    def __init__(
+        self, template_project_dir: Union[pathlib.Path, str], project_options: dict = None
+    ):
         self._project_options = project_options
 
-        if isinstance(template_project_dir, pathlib.Path):
+        if isinstance(template_project_dir, (pathlib.Path, str)):
             self._template_project_dir = str(template_project_dir)
         elif not isinstance(template_project_dir, str):
             raise TypeError(f"Incorrect type {type(template_project_dir)}.")

diff --git a/tests/python/unittest/test_crt.py b/tests/python/unittest/test_crt.py
@@ -248,7 +248,7 @@ def test_autotune():
     inputs = {"data": input_data}
 
     target = tvm.target.target.micro("host")
-    template_project_dir = os.path.join(tvm.micro.get_standalone_crt_dir(), "template", "host")
+    template_project_dir = pathlib.Path(tvm.micro.get_standalone_crt_dir()) / "template" / "host"
 
     pass_context = tvm.transform.PassContext(opt_level=3, config={"tir.disable_vectorize": True})
     with pass_context: