docstr: Called at the end of each stage to summarize progress
docstr_tok: ["Called","at","the","end","of","each","stage","to","summarize","progress"]
code: def on_stage_end(self, stage, stage_loss, epoch=None):
      "Called at the end of each stage to summarize progress"
      if self.sub_stage != SubStage.GENERATOR:
          return
  
      if stage == sb.Stage.TRAIN:
          self.train_loss = stage_loss
          g_loss = torch.tensor(self.metrics["G"])  # batch_size
          d_loss = torch.tensor(self.metrics["D"])  # batch_size
          print("Avg G loss: %.3f" % torch.mean(g_loss))
          print("Avg D loss: %.3f" % torch.mean(d_loss))
      else:
          if self.hparams.calculate_dnsmos_on_validation_set:
              stats = {
                  "SI-SNR": -stage_loss,
                  "pesq": 5 * self.pesq_metric.summarize("average")
                  - 0.5,
                  "stoi": -self.stoi_metric.summarize("average"),
                  "dnsmos": self.dnsmos_metric.summarize("average"),
              }
          else:
              stats = {
                  "SI-SNR": -stage_loss,
                  "pesq": 5 * self.pesq_metric.summarize("average")
                  - 0.5,
                  "stoi": -self.stoi_metric.summarize("average"),
              }
  
      if stage == sb.Stage.VALID:
          old_lr, new_lr = self.hparams.lr_annealing(
              5.0 - stats["pesq"]
          )
          sb.nnet.schedulers.update_learning_rate(
              self.g_optimizer, new_lr
          )
  
          if self.hparams.use_tensorboard:
              if (
                  self.hparams.calculate_dnsmos_on_validation_set
              ):  # Note: very time consuming........
                  valid_stats = {
                      "SI-SNR": -stage_loss,
                      "pesq": 5 * self.pesq_metric.summarize("average")
                      - 0.5,
                      "stoi": -self.stoi_metric.summarize("average"),
                      "dnsmos": self.dnsmos_metric.summarize("average"),
                  }
              else:
                  valid_stats = {
                      "SI-SNR": -stage_loss,
                      "pesq": 5 * self.pesq_metric.summarize("average")
                      - 0.5,
                      "stoi": -self.stoi_metric.summarize("average"),
                  }
  
              self.hparams.tensorboard_train_logger.log_stats(
                  {"lr": old_lr}, valid_stats
              )
          self.hparams.train_logger.log_stats(
              {"Epoch": epoch, "lr": old_lr},
              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 shutil
  import torch
  import torchaudio
  import speechbrain
  import numpy
  import json
  import pickle
  import requests
  import time
  from urllib.parse import urlparse, urljoin
  from srmrpy import srmr
  from pesq import pesq
  from enum import Enum, auto
  from hyperpyyaml import load_hyperpyyaml
  from speechbrain.utils.metric_stats import MetricStats
  from speechbrain.processing.features import spectral_magnitude
  from speechbrain.nnet.loss.stoi_loss import stoi_loss
  from speechbrain.utils.distributed import run_on_main
  from speechbrain.dataio.sampler import ReproducibleWeightedRandomSampler
label_desc: Use image node_id 13 for calling the MetricGanBrain obj's underlying member method code with example usage: obj.on_stage_end(stage, stage_loss, epoch) without return types