diff --git a/checkpoint/orbax/checkpoint/_src/path/step.py b/checkpoint/orbax/checkpoint/_src/path/step.py index 8caa160a8..fef0a6203 100644 --- a/checkpoint/orbax/checkpoint/_src/path/step.py +++ b/checkpoint/orbax/checkpoint/_src/path/step.py @@ -395,7 +395,7 @@ def _glob_step_paths(self, base_path: epath.PathLike) -> list[epath.Path]: return [ epath.Path(f'gs://{bucket_name}/{folder}') for folder in result.prefixes - if folder.startswith(os.path.join(path_prefix, self.step_prefix)) + if folder.startswith(os.path.join(path_prefix, step_prefix_with_underscore(self.step_prefix))) ] else: return list(