From cda59c287f90442004eeeb82e81789db26f0e196 Mon Sep 17 00:00:00 2001 From: Michael Lazos Date: Fri, 30 Jun 2023 17:36:23 +0000 Subject: [PATCH] Don't limit fusions with foreach --- torch/_inductor/scheduler.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/torch/_inductor/scheduler.py b/torch/_inductor/scheduler.py index ec402ddc30c3..47da9c88d0ba 100644 --- a/torch/_inductor/scheduler.py +++ b/torch/_inductor/scheduler.py @@ -1176,7 +1176,11 @@ def can_fuse(self, node1: BaseSchedulerNode, node2: BaseSchedulerNode): ): return False # heuristic not needed for correctness - if len(node1.get_nodes()) + len(node2.get_nodes()) > config.max_fusion_size: + if ( + not node1.is_foreach() + and not node2.is_foreach() + and len(node1.get_nodes()) + len(node2.get_nodes()) > config.max_fusion_size + ): return False # heuristic not needed for correctness if node1.get_names() & node2.recursive_predecessors: