Move the reference model (for train correctness checks) to cpu before starting tests

[davidberard98]

Stack from ghstack:

• -> #1275

On some OSS models we see CUDA OOM if we enable train correctness checks.

For certain models, we can prevent this OOM by copying the reference model to CPU before running the tests.