def METHOD_NAME(self, step: int, epoch: int, train_metrics: Dict) -> None:
    if len(jax.devices()) > 1:
        train_metrics = jax_utils.unreplicate(train_metrics)
    metrics = {"train/loss": train_metrics["loss"], "epoch": epoch}
    wandb.log(metrics, step=step + 1)