aggregate option cleanup.

PiperOrigin-RevId: 670598813
google-research · Sep 3, 2024 · 4d5716e · 4d5716e
1 parent 59a473c
commit 4d5716e
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/t5x/checkpoints.py b/t5x/checkpoints.py
@@ -2056,7 +2056,7 @@ def _construct_save_args(
   """Create SaveArgs for Orbax saving."""
   if param_info.name.split('.')[0] != 'target':
     dtype = None
-  return ocp.SaveArgs(aggregate=False, dtype=dtype)
+  return ocp.SaveArgs(dtype=dtype)
 
 
 def _construct_restore_args(