From e2280572cd380dedb01c65f0f36ecb8686e55332 Mon Sep 17 00:00:00 2001 From: Mridul Sahu Date: Wed, 19 Nov 2025 00:17:14 -0800 Subject: [PATCH] Always write array metadata if self._array_metadata_store is not None. PiperOrigin-RevId: 834158652 --- pathwaysutils/persistence/orbax_handler.py | 14 +++++++------- 1 file changed, 7 insertions(+), 7 deletions(-) diff --git a/pathwaysutils/persistence/orbax_handler.py b/pathwaysutils/persistence/orbax_handler.py index 850e2a5..539f834 100644 --- a/pathwaysutils/persistence/orbax_handler.py +++ b/pathwaysutils/persistence/orbax_handler.py @@ -128,14 +128,14 @@ async def serialize( ) arrays.append(v) - metadata_coroutine = None - if any_random_key: - if self._array_metadata_store is None: - raise ValueError( - "Array metadata store is not set with a checkpoint that requires" - f" it. Array metadata: {array_metadatas}" - ) + if any_random_key and self._array_metadata_store is None: + raise ValueError( + "Array metadata store is not set with a checkpoint that requires" + f" it. Array metadata: {array_metadatas}" + ) + metadata_coroutine = None + if self._array_metadata_store is not None: metadata_coroutine = self._array_metadata_store.write( checkpoint_dir=infos[0].parent_dir, array_metadatas=array_metadatas,