[AutoScheduler] Fix custom build func in PopenWorker (#8939)

* [AutoScheduler] Fix custom build func in PopenWorker * Add assertion
apache · Sep 8, 2021 · 332afae · 332afae
1 parent 1dc1707
commit 332afae
Showing 1 changed file with 4 additions and 5 deletions.
diff --git a/python/tvm/auto_scheduler/measure.py b/python/tvm/auto_scheduler/measure.py
@@ -661,10 +661,6 @@ def local_build_worker(args):
         The build result of this Builder thread.
     """
     inp, build_func, verbose = args
-    assert build_func == BuildFunc.name, (
-        "BuildFunc.name: " + BuildFunc.name + ", but args is: " + build_func
-    )
-    build_func = BuildFunc.build_func
 
     return _local_build_worker(inp, build_func, verbose)
 
@@ -693,6 +689,9 @@ def local_builder_build(inputs, timeout, n_parallel, build_func="default", verbo
     res : List[BuildResult]
         The build results of these MeasureInputs.
     """
+    assert build_func == BuildFunc.name, (
+        "BuildFunc.name: " + BuildFunc.name + ", but args is: " + build_func
+    )
     executor = PopenPoolExecutor(
         n_parallel, timeout, reset_global_scope, (AutotvmGlobalScope.current,)
     )
@@ -701,7 +700,7 @@ def local_builder_build(inputs, timeout, n_parallel, build_func="default", verbo
         [
             (
                 i.serialize(),
-                build_func,
+                BuildFunc.build_func,
                 verbose,
             )
             for i in inputs