Save StepMetadata in the Checkpointer.

PiperOrigin-RevId: 707670612
google-research · Dec 19, 2024 · 30962ec · 30962ec
1 parent 5f03619
commit 30962ec
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/t5x/utils.py b/t5x/utils.py
@@ -256,6 +256,7 @@ def save(
       path: str,
       item: train_state_lib.TrainState,
       force: bool = False,
+      custom: dict[str, Any] | None = None,
       state_transformation_fns: Sequence[
           checkpoints.SaveStateTransformationFn
       ] = (),
@@ -268,6 +269,7 @@ def save(
       path: path to save item to.
       item: a TrainState PyTree to save.
       force: unused.
+      custom: unused.
       state_transformation_fns: Transformations to apply, in order, to the state
         before writing.
       concurrent_gb: the approximate number of gigabytes of partitionable