Small update to keep elastic training working

Author: lukebaumann

MaxText/elastic_train.py

@@ -125,6 +125,7 @@ def elastic_handler(
     )
     with mesh:
       data_iterator, _ = create_data_iterator(config, mesh)
+      input_data_shardings = maxtext_utils.get_input_data_sharding(config, mesh)
 
       step, snapshot_jax_arrays, _ = elastic_manager.get_resharded_snapshot(mesh)
 
@@ -187,6 +188,7 @@ def elastic_handler(
       learning_rate_schedule,
       metric_logger,
       writer,
+      input_data_shardings,
   )
 
 
@@ -348,6 +350,7 @@ def train_loop(config, elastic_manager, recorder, state=None):
             learning_rate_schedule,
             metric_logger,
             writer,
+            input_data_shardings,
         ) = ret
 
       if step == start_step:
@@ -378,6 +381,7 @@ def train_loop(config, elastic_manager, recorder, state=None):
             learning_rate_schedule,
             metric_logger,
             writer,
+            input_data_shardings,
         ) = ret
 
   if checkpoint_manager is not None: