Fix metric accumulation (cast to float32)

Author: razvanmatisan

trainer.py

@@ -221,6 +221,7 @@ def loss_classifier(self, params, dropout_rng, batch: Batch, is_training: bool):
             "loss": (loss * batch_size, batch_size),
             "accuracy": (correct_pred.sum(), batch_size),
         }
+        metrics = jax.tree_map(lambda x: x.astype(jnp.float32), metrics)
 
         return loss, metrics
 
@@ -246,6 +247,7 @@ def loss_autoregressor(self, params, dropout_rng, batch: Batch, is_training: boo
 
         batch_size = batch.patches.shape[0]
         metrics = {"loss": (loss * batch_size, batch_size)}
+        metrics = jax.tree_map(lambda x: x.astype(jnp.float32), metrics)
 
         return loss, metrics