restore best weights option

[StefanBloemheuvel]

Hi creators of Talos,

First of all, amazing work on this package!
I would like to ask if there is a way to let the experiments restore the best weights of the best epoch, instead of saving the last epoch. I tried adding a callback, but that did not lead to success. I also tried the talos.utils.early_stopper, but it does not have the restore_best_weights option.

Kind regards!

Stefan