From 4134e0036ca9d84e0069957cacbfd845c9a8a425 Mon Sep 17 00:00:00 2001 From: Jez Ng Date: Thu, 22 Jun 2023 14:03:19 -0700 Subject: [PATCH] better implementation of pruning --- .../aitemplate/compiler/transform/transform_merge_view_ops.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/python/aitemplate/compiler/transform/transform_merge_view_ops.py b/python/aitemplate/compiler/transform/transform_merge_view_ops.py index 8e0b988be..cfb426df6 100644 --- a/python/aitemplate/compiler/transform/transform_merge_view_ops.py +++ b/python/aitemplate/compiler/transform/transform_merge_view_ops.py @@ -61,7 +61,8 @@ def _merge_view_ops_for(graph: List[Tensor], tensor: Tensor) -> List[Tensor]: graph.append(new_out_tensor) graph.remove(out_tensor) removed_ops.add(op) - dst_ops -= removed_ops + for op in removed_ops: + transform_utils.remove_view_op_from_sorted_graph(op) return graph