Enable to autotune maxtext extra_flags and maxtext_flags.

PiperOrigin-RevId: 837627814

Author: Google-ML-Automation

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