Fix usage of mask for computing total in Std metric.

PiperOrigin-RevId: 529696402

Author: Qwlouse

clu/metrics.py

@@ -815,8 +815,10 @@ def from_model_output(
     if mask is None:
       mask = jnp.ones(values.shape[0], dtype=jnp.int32)
     return cls(
-        total=values.sum(),
-        sum_of_squares=jnp.where(mask, values**2, jnp.zeros_like(values)).sum(),
+        total=jnp.where(mask, values, jnp.zeros_like(values)).sum(),
+        sum_of_squares=jnp.where(
+            mask, values**2, jnp.zeros_like(values)
+        ).sum(),
         count=mask.sum(),
     )