fix all instances where a T5X loss_fn/eval_fn is called to ensure we handle the return values correctly regardless of whether or not `weight_sum` is present as an auxiliary value.
fix all instances where a T5X loss_fn/eval_fn is called to ensure we handle the return values correctly regardless of whether or not `weight_sum` is present as an auxiliary value.

PiperOrigin-RevId: 440172420

google-research/prompt-tuning,google-research/prompt-tuning