Add an option to execute eval on cpu, regardless of training runs on TPU.

This will let users to benefit from TPU training, but avoid complex
eval metrics functions to be ported to TPU.

PiperOrigin-RevId: 196587755