From 1df6bb6d3080789e21b9eec8924949463593c757 Mon Sep 17 00:00:00 2001 From: Wuwei Lin Date: Tue, 14 Apr 2020 02:47:57 -0400 Subject: [PATCH] [TE][BuildModule] Fix import in dump pass ir (#5327) --- python/tvm/target/build_config.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/python/tvm/target/build_config.py b/python/tvm/target/build_config.py index 6a0dcf743a0d..8aae6be54a8b 100644 --- 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