Internal change.

PiperOrigin-RevId: 691549644
google · Oct 30, 2024 · 804e803 · 804e803
1 parent 712c7c7
commit 804e803
Showing 1 changed file with 5 additions and 2 deletions.
diff --git 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,
   )