class_docstr: The SpeechBrain class to train Non-Negative Factorization with Amortized Inference
docstr: 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."""
      # Compute/store important stats
      if stage == sb.Stage.TRAIN:
          self.train_loss = stage_loss
          self.train_stats = {
              "loss": self.train_loss,
          }
      # Summarize Valid statistics from the stage for record-keeping.
      elif stage == sb.Stage.VALID:
          valid_stats = {
              "loss": stage_loss,
          }
      # Perform end-of-iteration things, like annealing, logging, etc.
      if stage == sb.Stage.VALID:
          # The train_logger writes a summary to stdout and to the logfile.
          self.hparams.train_logger.log_stats(
              stats_meta={"epoch": epoch},
              train_stats=self.train_stats,
              valid_stats=valid_stats,
          )
          # Save the current checkpoint and delete previous checkpoints,
          self.checkpointer.save_and_keep_only(
              meta=valid_stats, min_keys=["loss"]
          )
  
file_path: speechbrain/recipes/ESC50/interpret/train_nmf.py
imports: import sys
  import torch
  import speechbrain
  from hyperpyyaml import load_hyperpyyaml
  from speechbrain.utils.distributed import run_on_main
  from esc50_prepare import prepare_esc50
  from train_l2i import dataio_prep
label_desc: Use image node_id 3 for calling the NMFBrain obj's underlying member method code with example usage: obj.on_stage_end(stage, stage_loss, epoch) without return types