diff --git a/python/tvm/relax/transform/ipc_allreduce_rewrite.py b/python/tvm/relax/transform/ipc_allreduce_rewrite.py index 8dc535020b30..3e7b005a6089 100644 --- a/python/tvm/relax/transform/ipc_allreduce_rewrite.py +++ b/python/tvm/relax/transform/ipc_allreduce_rewrite.py @@ -23,7 +23,6 @@ import tvm from tvm import relax from tvm.ir.module import IRModule -from tvm.relax.analysis import remove_all_unused from tvm.relax.expr import Expr, Var from tvm.relax.expr_functor import PyExprMutator, PyExprVisitor, mutator, visitor @@ -138,7 +137,6 @@ def transform(self) -> IRModule: for g_var, func in self.mod.functions_items(): if isinstance(func, relax.Function): updated_func = self.visit_expr(func) - updated_func = remove_all_unused(updated_func) self.builder_.update_func(g_var, updated_func) return self.builder_.get() diff --git a/python/tvm/relax/transform/lower_gpu_ipc_alloc_storage.py b/python/tvm/relax/transform/lower_gpu_ipc_alloc_storage.py index 0967e007563e..00081f92b197 100644 --- a/python/tvm/relax/transform/lower_gpu_ipc_alloc_storage.py +++ b/python/tvm/relax/transform/lower_gpu_ipc_alloc_storage.py @@ -21,7 +21,6 @@ import tvm from tvm import relax from tvm.ir.module import IRModule -from tvm.relax.analysis import remove_all_unused from tvm.relax.expr import Expr from tvm.relax.expr_functor import PyExprMutator, mutator @@ -49,7 +48,6 @@ def transform(self) -> IRModule: for g_var, func in self.mod.functions_items(): if isinstance(func, relax.Function): updated_func = self.visit_expr(func) - updated_func = remove_all_unused(updated_func) self.builder_.update_func(g_var, updated_func) return self.builder_.get()