GitLab

Use sparse_categorical_crossentropy for test as the loss object default does not work on tpustrategy +
the single task trainer already handles the reduction.

PiperOrigin-RevId: 367757677

PiperOrigin-RevId: 367564187