code: def on_stage_end(self, stage, stage_loss, epoch):
      if stage == sb.Stage.TRAIN:
          self.train_loss = stage_loss
      else:
          stage_stats = {"loss": stage_loss}
          max_keys = []
          min_keys = []
  
          if self.hparams.enhance_weight > 0:
              stage_stats["enhance"] = self.enh_metrics.summarize(
                  "average"
              )
              stage_stats["stoi"] = self.stoi_metrics.summarize(
                  "average"
              )
              stage_stats["pesq"] = self.pesq_metrics.summarize(
                  "average"
              )
              max_keys.extend(["pesq", "stoi"])
              if stage == sb.Stage.TEST:
                  stage_stats[
                      "csig"
                  ] = self.composite_metrics.summarize("csig")
                  stage_stats[
                      "cbak"
                  ] = self.composite_metrics.summarize("cbak")
                  stage_stats[
                      "covl"
                  ] = self.composite_metrics.summarize("covl")
                  max_keys.extend(["csig", "cbak", "covl"])
  
          if self.hparams.mimic_weight > 0:
              stage_stats["mimic"] = self.mimic_metrics.summarize(
                  "average"
              )
              min_keys.append("mimic")
  
          if self.hparams.ctc_weight > 0 or self.hparams.seq_weight > 0:
              err_rate = self.err_rate_metrics.summarize("error_rate")
              err_rate_type = self.hparams.target_type + "ER"
              stage_stats[err_rate_type] = err_rate
              min_keys.append(err_rate_type)
  
      if stage == sb.Stage.VALID:
          stats_meta = {"epoch": epoch}
          if hasattr(self.hparams, "lr_annealing"):
              old_lr, new_lr = self.hparams.lr_annealing(epoch - 1)
              sb.nnet.schedulers.update_learning_rate(
                  self.optimizer, new_lr
              )
              stats_meta["lr"] = old_lr
  
          self.hparams.train_logger.log_stats(
              stats_meta=stats_meta,
              train_stats={"loss": self.train_loss},
              valid_stats=stage_stats,
          )
          self.checkpointer.save_and_keep_only(
              meta=stage_stats,
              max_keys=max_keys,
              min_keys=min_keys,
              num_to_keep=self.hparams.checkpoint_avg,
          )
  
      elif stage == sb.Stage.TEST:
          self.hparams.train_logger.log_stats(
              stats_meta={
                  "Epoch loaded": self.hparams.epoch_counter.current
              },
              test_stats=stage_stats,
          )
          if if_main_process():
              with open(self.hparams.stats_file + ".txt", "w") as w:
                  if self.hparams.enhance_weight > 0:
                      w.write("\nstoi stats:\n")
                      self.stoi_metrics.write_stats(w)
                      w.write("\npesq stats:\n")
                      self.pesq_metrics.write_stats(w)
                      w.write("\ncomposite stats:\n")
                      self.composite_metrics.write_stats(w)
                  if self.hparams.mimic_weight > 0:
                      w.write("\nmimic stats:\n")
                      self.mimic_metrics.write_stats(w)
                  if self.hparams.seq_weight > 0:
                      self.err_rate_metrics.write_stats(w)
                  print("stats written to ", self.hparams.stats_file)
  
imports: import os
  import sys
  import torch
  import torchaudio
  import logging
  import speechbrain
  from pesq import pesq
  from pystoi import stoi
  from composite_eval import eval_composite
  from hyperpyyaml import load_hyperpyyaml
  from speechbrain.utils.data_utils import undo_padding
  from speechbrain.utils.distributed import run_on_main, if_main_process
label_desc: Use image node_id 7 for calling the MTLbrain obj's underlying member method code with example usage: obj.on_stage_end(stage, stage_loss, epoch) without return types