In case you want to continue from the same iteration, you would need to store the model, optimizer, and learning rate scheduler state_dicts as well as the current epoch and iteration. Assuming you want to get the same training batch, you could iterate the DataLoader in an “empty” loop until the appropriate iteration is reached (you could also seed the code properly so that the same “random” transformations are used, if needed).