fix registry

apache · Jan 21, 2021 · ec7ae58 · ec7ae58
1 parent 5537521
commit ec7ae58
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/python/tvm/auto_scheduler/workload_registry.py b/python/tvm/auto_scheduler/workload_registry.py
@@ -215,6 +215,9 @@ def serialize_workload_registry_entry(workload_key):
     """
     global WORKLOAD_FUNC_REGISTRY
 
+    if workload_key in WORKLOAD_FUNC_REGISTRY:
+        return (workload_key, WORKLOAD_FUNC_REGISTRY[workload_key])
+
     workload = json.loads(workload_key)
     name = workload[0]
     value = WORKLOAD_FUNC_REGISTRY[name]