def METHOD_NAME(
    self, step: int, epoch: int, batch_loss: float, batch_outputs: List[Dict[str, Any]]
) -> None:
    peak_gpu_mbs = peak_gpu_memory()
    if self.is_local_main_process:
        metrics = {
            "train/loss": batch_loss,
            "train/lr": self.training_engine.optimizer.param_groups[0]["lr"],
            "epoch": epoch,
        }
        metrics.update(
            {f"sys/worker{rank}_peak_gpu_mem": mbs for rank, mbs in peak_gpu_mbs.items()}
        )
        wandb.log(
            metrics,
            step=step + 1,
        )