Write tree metadata when checkpointing.

PiperOrigin-RevId: 577883417
google-research · Oct 30, 2023 · a9b35bc · a9b35bc
1 parent 44906f2
commit a9b35bc
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/vmoe/train/trainer.py b/vmoe/train/trainer.py
@@ -153,7 +153,9 @@ def create_checkpoint_manager(
       directory,
       {
           'state': orbax.checkpoint.AsyncCheckpointer(
-              orbax.checkpoint.PyTreeCheckpointHandler(),
+              orbax.checkpoint.PyTreeCheckpointHandler(
+                  write_tree_metadata=True,
+              ),
               timeout_secs=wait_seconds,
           ),
           'dataset_iterator': orbax.checkpoint.Checkpointer(