fix

hpcaitech · Aug 27, 2024 · 79274ef · 79274ef
1 parent ddfc5cc
commit 79274ef
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/colossalai/pipeline/stage_manager.py b/colossalai/pipeline/stage_manager.py
@@ -94,8 +94,8 @@ def get_stage_index(
             stage_indices.append([num_layers_per_stage_accumulated[stage], num_layers_per_stage_accumulated[stage + 1]])
             stage_indices.append(
                 [
-                    num_layers_per_stage_accumulated[num_stages - stage - 1],
-                    num_layers_per_stage_accumulated[num_stages - stage],
+                    num_layers_per_stage_accumulated[2 * num_stages - stage - 1],
+                    num_layers_per_stage_accumulated[2 * num_stages - stage],
                 ]
             )
             return stage_indices