From 804e8032b71463ebae0463fe6ab22aedb66b4ace Mon Sep 17 00:00:00 2001 From: Niket Kumar Bhumihar Date: Wed, 30 Oct 2024 14:14:57 -0700 Subject: [PATCH] Internal change. PiperOrigin-RevId: 691549644 --- checkpoint/orbax/checkpoint/path/step.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/checkpoint/orbax/checkpoint/path/step.py b/checkpoint/orbax/checkpoint/path/step.py index 2385c74f1..96f2e708d 100644 --- a/checkpoint/orbax/checkpoint/path/step.py +++ b/checkpoint/orbax/checkpoint/path/step.py @@ -185,9 +185,12 @@ def latest_step_metadata( root_path: epath.PathLike, name_format: NameFormat[MetadataT] ) -> Optional[MetadataT]: """Returns step.MetadataT of the latest step in `root_path`.""" - root_path = epath.Path(root_path) return max( - name_format.find_all(root_path), + sorted( + name_format.find_all(root_path), + key=lambda metadata: metadata.path.name, + reverse=True, + ), default=None, key=lambda metadata: metadata.step, )