Internal change

PiperOrigin-RevId: 681197359
google · Oct 1, 2024 · b722f4d · b722f4d
1 parent 9c162ab
commit b722f4d
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/flax/training/checkpoints.py b/flax/training/checkpoints.py
@@ -45,6 +45,7 @@
 from flax import config, core, errors, io, serialization, traverse_util
 from flax.training import orbax_utils
 
+
 _READ_CHECKPOINT_EVENT: str = '/jax/checkpoint/read/durations_sec'
 _WRITE_CHECKPOINT_EVENT: str = '/jax/checkpoint/write/durations_sec'
 
@@ -1136,6 +1137,7 @@ def restore_checkpoint(
     )
     return restored
 
+  # Legacy Flax checkpoint restoration.
   ckpt_size = io.getsize(ckpt_path)
   with io.GFile(ckpt_path, 'rb') as fp:
     if parallel and fp.seekable():