Auto pick checkpoint based on metrics

I’m training a model, and let’s say I evaluate after every epoch, I want the best model to be saved at the end of the training session based on a training metric I provide. Does PyTorch provide a wrapper of a kind?