Method stop

axlearn/common/checkpointer_orbax.py:614–616

See `BaseCheckpointer.stop` for details.

(self, *, has_exception: bool = False)

            target_structure=self._get_spec(step=step, state=state)["index"],
            validation=cfg.validation_type,
        )
        return step, restored_state

    def wait_until_finished(self):
        """See `BaseCheckpointer.wait_until_finished` docstring for details."""
        self._manager.wait_until_finished()