From d9648bee829f0008a82921f0aab8c7977b67261b Mon Sep 17 00:00:00 2001 From: Zhiyu Li Date: Mon, 16 Oct 2023 22:05:35 +0000 Subject: [PATCH] profile last 2 steps --- MaxText/train.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/MaxText/train.py b/MaxText/train.py index bfdf8e6d8..3a51d6a16 100644 --- 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)