feat(train): add load_best_on_end argument to save checkpoint

[martinkim0]

with this feature, the save checkpoint callback is feature-complete w.r.t. the current savebeststate callback.

needed before addressing #2568

[codecov]

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 88.83%. Comparing base (85540d1) to head (aaac7c4).
Report is 1 commits behind head on main.

❗ Current head aaac7c4 differs from pull request most recent head 2c496e5. Consider uploading reports for the commit 2c496e5 to get more accurate results

Additional details and impacted files

@@           Coverage Diff           @@
##             main    #2672   +/-   ##
=======================================
  Coverage   88.83%   88.83%           
=======================================
  Files         158      158           
  Lines       13235    13242    +7     
=======================================
+ Hits        11757    11764    +7     
  Misses       1478     1478           

Files	Coverage Δ
scvi/train/_callbacks.py	85.92% <100.00%> (+0.76%)	⬆️