No public description

PiperOrigin-RevId: 627629261
google-research · Apr 24, 2024 · c65605b · c65605b
1 parent 2bc6c0e
commit c65605b
Showing 1 changed file with 1 addition and 3 deletions.
diff --git a/vmoe/train/trainer.py b/vmoe/train/trainer.py
@@ -153,9 +153,7 @@ def create_checkpoint_manager(
       directory,
       {
           'state': orbax.checkpoint.AsyncCheckpointer(
-              orbax.checkpoint.PyTreeCheckpointHandler(
-                  write_tree_metadata=True,
-              ),
+              orbax.checkpoint.PyTreeCheckpointHandler(),
               timeout_secs=wait_seconds,
           ),
           'dataset_iterator': orbax.checkpoint.Checkpointer(