Merge branch 'main' of github.com:Accelergy-Project/timeloop-python

Accelergy-Project · Nov 13, 2024 · 12094eb · 12094eb
2 parents 5c285cd + b366f18
commit 12094eb
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/pytimeloop/fastfusion/mapper/per_einsum_mapper.py b/pytimeloop/fastfusion/mapper/per_einsum_mapper.py
@@ -249,9 +249,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,