Fix arg input in RemoteIteratorWrapper

aireenmei · aireenmei · commit d5f2fd6fc2ae · 2026-05-09T02:00:32.000Z
diff --git a/src/maxtext/input_pipeline/tfds_data_processing.py b/src/maxtext/input_pipeline/tfds_data_processing.py
@@ -305,6 +305,7 @@ def make_tfds_eval_iterator(
         use_dpo=config.use_dpo,
         hf_access_token=config.hf_access_token,
     )
+    global_shape = (config.global_batch_size_to_load_eval, config.max_target_length)
     return multihost_dataloading.RemoteIteratorWrapper(
-        get_ds_fn, preprocessing_fn, config, global_mesh, checkpoint_path=config.checkpoint_dir
+        get_ds_fn, preprocessing_fn, global_mesh, global_shape, checkpoint_path=config.checkpoint_dir
     )

Original file line number	Diff line number	Diff line change
`@@ -305,6 +305,7 @@ def make_tfds_eval_iterator(`
`305`	`305`	`use_dpo=config.use_dpo,`
`306`	`306`	`hf_access_token=config.hf_access_token,`
`307`	`307`	`)`
	`308`	`+ global_shape = (config.global_batch_size_to_load_eval, config.max_target_length)`
`308`	`309`	`return multihost_dataloading.RemoteIteratorWrapper(`
`309`		`- get_ds_fn, preprocessing_fn, config, global_mesh, checkpoint_path=config.checkpoint_dir`
	`310`	`+ get_ds_fn, preprocessing_fn, global_mesh, global_shape, checkpoint_path=config.checkpoint_dir`
`310`	`311`	`)`