Fix bug in relax.vm.build to pass target argument. (#91)

Co-authored-by: Prakalp Srivastava <prakalp@octoml.ai>
tlc-pack · Feb 8, 2023 · 6ff0632 · 6ff0632
1 parent 9355371
commit 6ff0632
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 1 deletion.
diff --git a/python/tvm/relax/vm.py b/python/tvm/relax/vm.py
@@ -186,7 +186,7 @@ def foo(x: Tensor[(3, 4), "float32"], y: Tensor[(3, 4), "float32"]):
 
     # split primfunc and relax function
     rx_mod, tir_mod = _split_tir_relax(new_mod)
-    lib = tvm.build(tir_mod, target)
+    lib = tvm.build(tir_mod, target=target)
     ex = _ffi_api.VMCodeGen(rx_mod)
     return ex, lib
 

diff --git a/tests/python/relax/test_autotir_integration.py b/tests/python/relax/test_autotir_integration.py
@@ -29,6 +29,7 @@
 from tvm.meta_schedule.integration import extract_task_from_relax
 from tvm import transform
 import time
+import pytest
 
 # Test case with dynamic shape.
 # Tuning with dynamic shape is not supported yet.
@@ -115,6 +116,7 @@ def print_results(self) -> None:
         print("\n".join([str(r) for r in self.records]))
 
 
+@pytest.mark.parametrize("dev", ["cpu"])
 def test_class_irmodule(dev: str):
     @tvm.script.ir_module
     class InputModule: