How to save a cehckpoint after each epoch - dataloader states

I need to save a checkpoint after each epoch with I see how this can be done with save_steps for after X number of steps but I need this to be done after each epoch, any ideas ?

One more question, if I save the cehckpoint after each X steps, when I want to retrain, dataloader states are not the same, is there an idea how to solve this issue? thanks