From b366f181f6a004e1400e1d979f5d5130798bc188 Mon Sep 17 00:00:00 2001
From: Michael Gilbert <gilbertm@mit.edu>
Date: Wed, 13 Nov 2024 15:29:04 -0500
Subject: [PATCH] Fix bug with fused-unfused exploration

---
 pytimeloop/fastfusion/mapper/per_einsum_mapper.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/pytimeloop/fastfusion/mapper/per_einsum_mapper.py b/pytimeloop/fastfusion/mapper/per_einsum_mapper.py
index 79fff7b..341364a 100644
--- a/pytimeloop/fastfusion/mapper/per_einsum_mapper.py
+++ b/pytimeloop/fastfusion/mapper/per_einsum_mapper.py
@@ -251,9 +251,9 @@ def get_top_loop_jobs(
         logfunc(f"Allowed top-level loop ranks: {top_level_ranks}")
         for partial_mapping in explore_fused_unfused(mapping,
                                                      intermediate_tensors):
-            for partial_mapping in make_top_loops(mapping,
-                                                top_level_ranks,
-                                                logfunc):
+            for partial_mapping in make_top_loops(partial_mapping,
+                                                  top_level_ranks,
+                                                  logfunc):
                 for partial_mapping in place_glb_level(
                     partial_mapping,
                     intermediate_tensors,