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')
model.train()
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?