Modify Welford metric to return mean value.

PiperOrigin-RevId: 640835234

flax/nnx/nnx/training/metrics.py

@@ -128,7 +128,7 @@ def compute(self):
     standard_deviation = variance**0.5
     sem = standard_deviation / (self.count**0.5)
     return Statistics(
-        mean=self.mean,
+        mean=self.mean.value,
         standard_error_of_mean=sem,
         standard_deviation=standard_deviation,
     )