Performance improvement for distillation of DNNs to trees.

* Move the DNN network to the training-only graph so it is not in the eval if we're doing distillation. This makes the inference much faster because we no longer run the DNN.
* Simplify use_core_version logic. Now we always return model_fn for contrib and estimator_spec for core.

PiperOrigin-RevId: 216932818