Update checkpoint path in load_models function to use the latest training run

init commit