[FIX] fix transform params without GPUs (octoml#219)

Lunderberg · May 23, 2023 · 110c6d3 · 110c6d3
1 parent 68be032
commit 110c6d3
Showing 1 changed file with 4 additions and 2 deletions.
diff --git a/mlc_llm/utils.py b/mlc_llm/utils.py
@@ -145,8 +145,10 @@ def set_item(i, value):
         res[i] = tvm.nd.array(value, device=tvm.cpu())
         return tvm.nd.empty((1,), device=device)
 
-    with tvm.target.Target(target):
-        mod_transform = tvm.tir.transform.DefaultGPUSchedule()(mod_transform)
+    if target.kind.name != "llvm":
+        with tvm.target.Target(target):
+            mod_transform = tvm.tir.transform.DefaultGPUSchedule()(mod_transform)
+
     ex = relax.build(mod_transform, target=target)
     vm = relax.vm.VirtualMachine(ex, device)
     vm[transform_func_name]()