def on_stage_end(self, stage, stage_loss, epoch):
        """Gets called at the end of a epoch."""

        # Compute/store important stats
        stage_stats = {"loss": stage_loss}
        if stage == sb.Stage.TRAIN:
            print(f'[rank {torch.distributed.get_rank()}] Epoch {epoch} training loss: {stage_loss:.4f}')
            self.train_stats = stage_stats
        else:
            stage_stats["CER"] = self.cer_metric.summarize("error_rate")

        # Perform end-of-iteration things, like annealing, logging, etc.
        if stage == sb.Stage.VALID:

            lr = self.hparams.noam_annealing.current_lr
            steps = self.optimizer_step
            optimizer = self.optimizer.__class__.__name__

            epoch_stats = {
                "epoch": epoch,
                "lr": lr,
                "steps": steps,
                "optimizer": optimizer,
            }

            self.hparams.train_logger.log_stats(
                stats_meta=epoch_stats,
                train_stats=self.train_stats,
                valid_stats=stage_stats,
            )
            print(f'[rank {torch.distributed.get_rank()}] Epoch {epoch} CER {stage_stats["CER"]:.4f}')
            self.checkpointer.save_and_keep_only(
                meta={"CER": stage_stats["CER"], "epoch": epoch},
                min_keys=["CER"],
                num_to_keep=self.hparams.avg_checkpoints,
                ckpt_predicate=lambda ckpt: ckpt.meta.get("CER")!=-0.1
            )