From 9982f677d714ea7808b2e8c0aa04496e542dfd12 Mon Sep 17 00:00:00 2001 From: Orbax Authors Date: Mon, 15 Jun 2026 22:04:25 -0700 Subject: [PATCH] Roll forward of shape reporting with a fix for UnparsedFlagAccessError PiperOrigin-RevId: 932865483 --- .../checkpoint/_src/handlers/base_pytree_checkpoint_handler.py | 1 + 1 file changed, 1 insertion(+) diff --git a/checkpoint/orbax/checkpoint/_src/handlers/base_pytree_checkpoint_handler.py b/checkpoint/orbax/checkpoint/_src/handlers/base_pytree_checkpoint_handler.py index 98aea2be5..f25e2f5fa 100644 --- a/checkpoint/orbax/checkpoint/_src/handlers/base_pytree_checkpoint_handler.py +++ b/checkpoint/orbax/checkpoint/_src/handlers/base_pytree_checkpoint_handler.py @@ -26,6 +26,7 @@ import dataclasses import functools import json +import os import threading import time from typing import Any, List, Mapping, Optional, Tuple, Union