Situation: I train for a while, then want to save exactly the current train state to disk, and exit. Then I want to continue training later. And it should be exactly the sam