Copybara import of the project:

--
40e056d by NuojCheng <[email protected]>:

add rampup batch size

COPYBARA_INTEGRATE_REVIEW=AI-Hypercomputer#2535 from AI-Hypercomputer:chengnuojin-rampup-batch 40e056d
PiperOrigin-RevId: 827037473

Author: NuojCheng

tests/data_loader_test.py

@@ -51,15 +51,6 @@ def setUp(self):
 
   def get_test_config(self, reuse_example_batch, **kwargs):
     """Generate config for tests"""
-    args = {
-        "run_name": "test",
-        "mesh_axes": ["data"],
-        "logical_axis_rules": [["batch", "data"]],
-        "data_sharding": ["data"],
-        "enable_checkpointing": False,
-        "reuse_example_batch": reuse_example_batch,
-    }
-    args.update(kwargs)
     return pyconfig.initialize(
       [None, get_test_config_path()],
       per_device_batch_size=1,
@@ -164,4 +155,4 @@ def test_rampup_data_loader(self):
 
 
 if __name__ == "__main__":
-  unittest.main()
+  unittest.main()