Restarting training from a checkpoint

I’m following through the Training a diffusers model notebook and I’ve adapted it and run 50 epochs on my own dataset. I’m running on Colab so I need to do training in small-ish bursts, so I’d like to set the notebook up to start training from the last checkpoint.

Reading through the code I saw how to load a pretrained model and that it needs to be set from eval mode back to train, which I’ve done with:

model = UNet2DModel.from_pretrained(previous_checkpoint_dir + '/unet')

I then ran the rest of the notebook code as-is (apart of some code to find the last checkpoint) to restart the training. Is there any more to it than that?