Returns the path of checkpoint at `step` under `base_dir`.
(base_dir: str, *, step: int)
| 120 | |
| 121 | |
| 122 | def build_step_dir(base_dir: str, *, step: int) -> str: |
| 123 | """Returns the path of checkpoint at `step` under `base_dir`.""" |
| 124 | return os.path.join(base_dir, f"{STEP_PREFIX}_{step:0{STEP_NUM_DIGITS}d}") |
| 125 | |
| 126 | |
| 127 | def check_state_structure( |