Add logging to track deprecated codepaths.

PiperOrigin-RevId: 652597860
google · Aug 12, 2024 · ec9f3de · ec9f3de
1 parent 8e0228e
commit ec9f3de
Showing 1 changed file with 1 addition and 4 deletions.
diff --git a/flax/training/orbax_utils.py b/flax/training/orbax_utils.py
@@ -33,10 +33,7 @@ def is_multi_device_array(value: Any) -> bool:
 
 
 def save_args_from_target(target: Any) -> Any:
-  return jax.tree_util.tree_map(
-    lambda x: ocp.SaveArgs(aggregate=not is_multi_device_array(x)),
-    target,
-  )
+  return jax.tree_util.tree_map(lambda _: ocp.SaveArgs(), target)
 
 
 def maybe_construct_transformations(