Make `select_checkpoint_from_menu`'s behavior more generic over filenames

Allow leaving certain parts of the model to be uninitialized from checkpoints

Small changes and fixes

Update train script after clean up

Reorganization and cleanups

Initial commit