Update the `loss_fn` and `eval_fn` of the T5X models to remove `weight_sum` from the return values.
Update the `loss_fn` and `eval_fn` of the T5X models to remove `weight_sum` from the return values.

PiperOrigin-RevId: 445292577

google-research/t5x