Revert "profile last 2 steps"

This reverts commit f7fb49a.
google · Nov 14, 2023 · 00c27ea · 00c27ea
1 parent f7fb49a
commit 00c27ea
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/MaxText/train.py b/MaxText/train.py
@@ -295,7 +295,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 == config.steps - 2:
+    if step == 0:
       max_utils.activate_profiler(config)
 
   max_utils.deactivate_profiler(config)