Method restore

axlearn/common/checkpointer_test.py:1433–1441

()

                    storage.wait_until_finished()

                asyncio.run(save())

                async def restore():
                    # Successfully restores with different dtypes.
                    restored_state = storage.restore_from_dir(
                        step,
                        state=make_state(float_dtype=restore_floats_as),
                        ckpt_dir=final_dir,
                        validation=CheckpointValidationType.EXACT_UP_TO_DTYPE,
                    )
                    return restored_state