From 624c63c470e528a196ca1044d4cce23dceed30bf Mon Sep 17 00:00:00 2001 From: Colin Gaffney Date: Thu, 23 May 2024 13:51:04 -0700 Subject: [PATCH] Stop writing msgpack file for new checkpoints and update empty nodes handling so that it no longer depends on this file. PiperOrigin-RevId: 636665054 --- flax/training/checkpoints.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/flax/training/checkpoints.py b/flax/training/checkpoints.py index 747ba63431..3d48731305 100644 --- a/flax/training/checkpoints.py +++ b/flax/training/checkpoints.py @@ -76,7 +76,7 @@ # Orbax main checkpoint file name. ORBAX_CKPT_FILENAME = 'checkpoint' -ORBAX_MANIFEST_OCDBT = 'manifest.ocdbt' +ORBAX_METADATA_FILENAME = '_METADATA' PyTree = Any @@ -123,7 +123,7 @@ def _safe_remove(path: str): def _is_orbax_checkpoint(path: str) -> bool: return io.exists(os.path.join(path, ORBAX_CKPT_FILENAME)) or io.exists( - os.path.join(path, ORBAX_MANIFEST_OCDBT) + os.path.join(path, ORBAX_METADATA_FILENAME) )