docstr: Gets called at the end of a epoch.
docstr_tok: ["Gets","called","at","the","end","of","a","epoch","."]
code: def on_stage_end(self, stage, stage_loss, epoch):
      """Gets called at the end of a epoch."""
      # Compute/store important stats
      stage_stats = {"si-snr": stage_loss}
      if stage == sb.Stage.TRAIN:
          self.train_stats = stage_stats
  
      # Perform end-of-iteration things, like annealing, logging, etc.
      if stage == sb.Stage.VALID:
          # Learning rate annealing
          if isinstance(
              self.hparams.lr_scheduler, schedulers.ReduceLROnPlateau
          ):
              current_lr, next_lr = self.hparams.lr_scheduler(
                  [self.optimizer], epoch, stage_loss
              )
              schedulers.update_learning_rate(self.optimizer, next_lr)
          else:
              # if we do not use the reducelronplateau, we do not change the lr
              current_lr = self.hparams.optimizer.optim.param_groups[0][
                  "lr"
              ]
  
          self.hparams.train_logger.log_stats(
              stats_meta={"epoch": epoch, "lr": current_lr},
              train_stats=self.train_stats,
              valid_stats=stage_stats,
          )
          if (
              hasattr(self.hparams, "save_all_checkpoints")
              and self.hparams.save_all_checkpoints
          ):
              self.checkpointer.save_checkpoint(
                  meta={"si-snr": stage_stats["si-snr"]}
              )
          else:
              self.checkpointer.save_and_keep_only(
                  meta={"si-snr": stage_stats["si-snr"]},
                  min_keys=["si-snr"],
              )
      elif stage == sb.Stage.TEST:
          self.hparams.train_logger.log_stats(
              stats_meta={
                  "Epoch loaded": self.hparams.epoch_counter.current
              },
              test_stats=stage_stats,
          )
  
imports: import os
  import sys
  import torch
  import torch.nn.functional
  import torchaudio
  import speechbrain
  import speechbrain.nnet.schedulers
  from speechbrain.utils.distributed import run_on_main
  from hyperpyyaml import load_hyperpyyaml
  import numpy
  from tqdm import tqdm
  import csv
  import logging
  from speechbrain.core import AMPConfig
label_desc: Use image node_id 5 for calling the Separation obj's underlying member method code with example usage: obj.on_stage_end(stage, stage_loss, epoch) without return types