Method wait_until_finished

axlearn/common/checkpointer_orbax.py:610–612

See `BaseCheckpointer.wait_until_finished` docstring for details.

(self)

        # Validate ckpt structure.
        check_state_structure(
            restored_index,
            target_structure=self._get_spec(step=step, state=state)["index"],
            validation=cfg.validation_type,
        )
        return step, restored_state