Hi @mralexis, as described in the docs 1, load_best_model_at_end  works in conjunction with metric_for_best_model.
Under the hood, the Trainer keeps track of the best checkpoint seen to date via the Trainer.state.best_model_checkpoint attribute (see here 2) which is then used by load_best_model_at_end here 3.