From 2a71e665e5919aec49bb5c17aaac0e7fec6b17aa Mon Sep 17 00:00:00 2001 From: maxtext authors Date: Thu, 27 Nov 2025 15:37:12 -0800 Subject: [PATCH] Enable to autotune maxtext `extra_flags` and `maxtext_flags`. PiperOrigin-RevId: 837627814 --- src/MaxText/train.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/MaxText/train.py b/src/MaxText/train.py index fb5fcebdb..bebd737e5 100644 --- a/src/MaxText/train.py +++ b/src/MaxText/train.py @@ -412,7 +412,7 @@ def train_loop(config, recorder, state=None): max_utils.print_compiled_memory_stats(compiled_stats) start_step = get_first_step(state) # this is the start_step for training - prof = profiler.Profiler(config, offset_step=start_step) + prof = profiler.Profiler(config, offset_step=start_step, ) metric_logger = MetricLogger(config=config, learning_rate_schedule=learning_rate_schedule) # Write train config params, num model params, and XLA flags to tensorboard