docstr: Gets called at the end of an epoch.
docstr_tok: ["Gets","called","at","the","end","of","an","epoch","."]
code: def on_stage_end(self, stage, stage_loss, epoch=None):
      """Gets called at the end of an epoch."""
      if stage == sb.Stage.TRAIN:
          self.train_loss = stage_loss
          self.train_stats = {  # "loss": self.loss_metric.scores,
              "loss_d1": self.loss_metric_d1.scores,
              "loss_d2": self.loss_metric_d2.scores,
              "loss_g3": self.loss_metric_g3.scores,
          }
      else:
          stats = {
              "loss": stage_loss,
              "pesq": self.pesq_metric.summarize("average"),
              "stoi": -self.stoi_metric.summarize("average"),
          }
  
      if stage == sb.Stage.VALID:
          if self.hparams.use_tensorboard:
              valid_stats = {
                  # "loss": self.loss_metric.scores,
                  "loss_d1": self.loss_metric_d1.scores,
                  "loss_d2": self.loss_metric_d2.scores,
                  "loss_g3": self.loss_metric_g3.scores,
                  "stoi": self.stoi_metric.scores,
                  "pesq": self.pesq_metric.scores,
              }
              self.hparams.tensorboard_train_logger.log_stats(
                  {"Epoch": epoch}, self.train_stats, valid_stats
              )
          self.hparams.train_logger.log_stats(
              {"Epoch": epoch},
              train_stats={"loss": self.train_loss},
              valid_stats=stats,
          )
          self.checkpointer.save_and_keep_only(
              meta=stats, max_keys=["pesq"]
          )
  
      if stage == sb.Stage.TEST:
          self.hparams.train_logger.log_stats(
              {"Epoch loaded": self.hparams.epoch_counter.current},
              test_stats=stats,
          )
  
imports: import os
  import sys
  import torch
  import torchaudio
  import speechbrain
  from pesq import pesq
  from hyperpyyaml import load_hyperpyyaml
  from speechbrain.utils.metric_stats import MetricStats
  from speechbrain.nnet.loss.stoi_loss import stoi_loss
  from speechbrain.utils.distributed import run_on_main
label_desc: Use image node_id 11 for calling the SEBrain obj's underlying member method code with example usage: obj.on_stage_end(stage, stage_loss, epoch) without return types