address comments

apache · Oct 1, 2021 · 9f293d5 · 9f293d5
1 parent 2d80cdc
commit 9f293d5
Showing 1 changed file with 6 additions and 8 deletions.
diff --git a/python/tvm/meta_schedule/runner/local_runner.py b/python/tvm/meta_schedule/runner/local_runner.py
@@ -238,7 +238,9 @@ def _check(
             get_global_func_with_default_on_worker(name=f_alloc_argument, default=None)
             get_global_func_with_default_on_worker(name=f_run_evaluator, default=None)
             get_global_func_with_default_on_worker(name=f_cleanup, default=None)
-            tvm.get_global_func("tvm.contrib.random.random_fill", True)
+            get_global_func_with_default_on_worker(
+                name="tvm.contrib.random.random_fill", default=None
+            )
 
         value = self.pool.submit(
             _check,
@@ -319,13 +321,9 @@ def default_alloc_argument(
     repeated_args: List[T_ARGUMENT_LIST]
         The allocation args
     """
-    try:
-        f_random_fill = tvm.get_global_func("tvm.contrib.random.random_fill", True)
-    except AttributeError as error:
-        raise AttributeError(
-            'Unable to find function "tvm.contrib.random.random_fill" on local runner. '
-            "Please make sure USE_RANDOM is turned ON in the config.cmake."
-        ) from error
+    f_random_fill = get_global_func_with_default_on_worker(
+        name="tvm.contrib.random.random_fill", default=None
+    )
     return alloc_argument_common(f_random_fill, device, args_info, alloc_repeat)