[TE][BuildModule] Fix import in dump pass ir (apache#5327)

siju-samuel · Apr 14, 2020 · 1df6bb6 · 1df6bb6
1 parent b7545eb
commit 1df6bb6
Showing 1 changed file with 5 additions and 5 deletions.
diff --git a/python/tvm/target/build_config.py b/python/tvm/target/build_config.py
@@ -59,15 +59,15 @@ def dump(*args, **kwargs):
 
     def decorate_irpass(self):
         """decorate ir_pass and ScheduleOps"""
-        self._old_sgpass = schedule.ScheduleOps
-        schedule.ScheduleOps = self.decorate(schedule.ScheduleOps)
-        vset = vars(ir_pass)
+        self._old_sgpass = tvm.te.schedule.ScheduleOps
+        tvm.te.schedule.ScheduleOps = self.decorate(tvm.te.schedule.ScheduleOps)
+        vset = vars(tvm.tir.ir_pass)
         k = v = 0
         def recover():
             vset[k] = v
         for k, v in vset.items():
             self._recover_list.append(recover)
-            vset[k] = self.decorate(v) if isinstance(v, Function) else v
+            vset[k] = self.decorate(v) if isinstance(v, tvm.runtime.PackedFunc) else v
 
     def decorate_custompass(self, custom_pass):
         """decorate given list of custom passes, and return decorated passes"""
@@ -93,7 +93,7 @@ def exit(self):
         # recover decorated functions
         for f in self._recover_list:
             f()
-        schedule.ScheduleOps = self._old_sgpass
+        tvm.te.schedule.ScheduleOps = self._old_sgpass
         DumpIR.scope_level -= 1