* Load checkpoint on CPU instead of on GPU

* refactor: simplify code

* Cleanup

* Update train.py
Co-authored-by: Glenn Jocher <glenn.jocher@ultralytics.com>