Change `update_state` call of sum-based metrics in `training_utils` to accept both the label and the model output.
Change `update_state` call of sum-based metrics in `training_utils` to accept both the label and the model output.

PiperOrigin-RevId: 352808820
Change-Id: I70e94512c03b97e541b10143b4ac43cc07a3d9b2

google-research/federated,google-research/federated