better implementation of pruning

facebookincubator · Jul 11, 2023 · 4134e00 · 4134e00
1 parent 64b1573
commit 4134e00
Showing 1 changed file with 2 additions and 1 deletion.
diff --git 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