How to save checkpoints?

[joandies]

I've seen that the function train has an input "checkpoint_to_load=None", I would like to know how to adapt the code so that after each epoch it saves its weights. Because I've calculated the total time of my training and it is around 48h, so I want to save my weights periodically.

def train(get_batches, data_shape, checkpoint_to_load=None):

Sorry if it's a stupid question, I'm a starter.