* Load checkpoint on CPU instead of on GPU * refactor: simplify code * Cleanup * Update train.py Co-authored-by: Glenn Jocher <glenn.jocher@ultralytics.com>