From 2130b1b65488dadecab9ae444b06e9f99cee2bd3 Mon Sep 17 00:00:00 2001
From: Alexander Ku <alexku@google.com>
Date: Mon, 18 Aug 2025 13:39:50 -0700
Subject: [PATCH] Fix checkpoint loading error due to Orbax update.

PiperOrigin-RevId: 796557269
---
 gemma/gm/ckpts/_checkpoint.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/gemma/gm/ckpts/_checkpoint.py b/gemma/gm/ckpts/_checkpoint.py
index 2cca3e08..8fab7b2e 100644
--- a/gemma/gm/ckpts/_checkpoint.py
+++ b/gemma/gm/ckpts/_checkpoint.py
@@ -442,7 +442,7 @@ def _get_metadata_and_path(
       metadata = ckpt.metadata(path)
     else:
       raise
-  metadata = dict(metadata)  # Normalize metadata
+  metadata = dict(metadata.item_metadata)  # Normalize metadata
   return metadata, path