Nice! But please add a store_best_model flag that defaults to False.