profile last 2 steps

google · Nov 22, 2023 · d9648be · d9648be
1 parent 07dc6ce
commit d9648be
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/MaxText/train.py b/MaxText/train.py
@@ -312,7 +312,7 @@ def train_loop(config, state=None):
 
     # Start profiling at end of first step to avoid compilation.
     # Move before for loop to include.
-    if step == 0:
+    if step == config.steps - 2:
       max_utils.activate_profiler(config)
 
   max_utils.deactivate_profiler(config)