Never lose your progress: Checkpointing with Orbax

Picture this. You have spent six hours training a model. The loss curve looks beautiful, accuracy is climbing, and you are one epoch away from a result worth writing home about. Then the power goes ou