class_docstr: Class that manages the training loop. See speechbrain.core.Brain.
  
  Arguments
  ---------
  modules : dict of str:torch.nn.Module pairs
      These modules are passed to the optimizer by default if they have
      trainable parameters, and will have ``train()``/``eval()`` called on them.
  opt_class : torch.optim class
      A torch optimizer constructor that takes only the list of
      parameters (e.g. a lambda or partial function definition). By default,
      this will be passed all modules in ``modules`` at the
      beginning of the ``fit()`` method. This behavior can be changed
      by overriding the ``configure_optimizers()`` method.
  hparams : dict
      Each key:value pair should consist of a string key and a hyperparameter
      that is used within the overridden methods. These will
      be accessible via an ``hparams`` attribute, using "dot" notation:
      e.g., self.hparams.model(x).
  run_opts : dict
      A set of options to change the runtime environment.
class_docstr_tok: ["Class","that","manages","the","training","loop",".","See","speechbrain.core.Brain",".","Arguments","--","--","--","--","-","modules",":","dict","of","str",":","torch.nn.Module","pairs","These","modules","are","passed","to","the","optimizer","by","default","if","they","have","trainable","parameters",",","and","will","have","``","train","(",")","``","\/","``","eval","(",")","``","called","on","them",".","opt_class",":","torch.optim","class","A","torch","optimizer","constructor","that","takes","only","the","list","of","parameters","(","e.g",".","a","lambda","or","partial","function","definition",")",".","By","default",",","this","will","be","passed","all","modules","in","``","modules","``","at","the","beginning","of","the","``","fit","(",")","``","method",".","This","behavior","can","be","changed","by","overriding","the","``","configure_optimizers","(",")","``","method",".","hparams",":","dict","Each","key",":","value","pair","should","consist","of","a","string","key","and","a","hyperparameter","that","is","used","within","the","overridden","methods",".","These","will","be","accessible","via","an","``","hparams","``","attribute",",","using","``","dot","''","notation",":","e.g.",",","self.hparams.model","(","x",")",".","run_opts",":","dict","A","set","of","options","to","change","the","runtime","environment","."]
docstr: Gets called at the end of an epoch.
  
  Arguments
  ---------
  stage : sb.Stage
      One of sb.Stage.TRAIN, sb.Stage.VALID, sb.Stage.TEST
  stage_loss : float
      The average loss for all of the data processed in this stage.
  epoch : int
      The currently-starting epoch. This is passed
      `None` during the test stage.
docstr_tok: ["Gets","called","at","the","end","of","an","epoch",".","Arguments","--","--","--","--","-","stage",":","sb.Stage","One","of","sb.Stage.TRAIN",",","sb.Stage.VALID",",","sb.Stage.TEST","stage_loss",":","float","The","average","loss","for","all","of","the","data","processed","in","this","stage",".","epoch",":","int","The","currently-starting","epoch",".","This","is","passed","`","None","`","during","the","test","stage","."]
code: def on_stage_end(self, stage, stage_loss, epoch=None):
      """Gets called at the end of an epoch.
  
      Arguments
      ---------
      stage : sb.Stage
          One of sb.Stage.TRAIN, sb.Stage.VALID, sb.Stage.TEST
      stage_loss : float
          The average loss for all of the data processed in this stage.
      epoch : int
          The currently-starting epoch. This is passed
          `None` during the test stage.
      """
  
      # Store the train loss until the validation stage.
      if stage == sb.Stage.TRAIN:
          self.train_loss = stage_loss
  
      # Summarize the statistics from the stage for record-keeping.
      else:
          stats = {"loss": stage_loss}
  
      # At the end of validation...
      if stage == sb.Stage.VALID:
          # The train_logger writes a summary to stdout and to the logfile.
          lr = self.optimizer.param_groups[0]["lr"]
          self.hparams.train_logger.log_stats(
              {"Epoch": epoch, "lr": lr},
              train_stats={"loss": self.train_loss},
              valid_stats=stats,
          )
  
          # Save the current checkpoint and delete previous checkpoints,
          self.checkpointer.save_and_keep_only(
              meta=stats, min_keys=["loss"]
          )
  
      # We also write statistics about test data to stdout and to the logfile.
      if stage == sb.Stage.TEST:
          self.hparams.train_logger.log_stats(
              {"Epoch loaded": self.hparams.epoch_counter.current},
              test_stats=stats,
          )
  
      if (
          stage == sb.Stage.TRAIN
          and self.hparams.enable_reference_samples
      ):
          self.generate_reference_samples(self.reference_batch)
  
      if (
          stage != sb.Stage.TRAIN
          and epoch is not None
          and epoch % self.hparams.samples_interval == 0
      ):
          labels, samples, samples_denorm, wav = self.generate_samples()
          samples_rec, wav_rec = None, None
          data = {
              "labels": labels,
              "samples": samples,
              "samples_denorm": samples_denorm,
              "wav": wav,
          }
          if self.diffusion_mode == DiffusionMode.LATENT:
              samples_rec, wav_rec = self.generate_rec_samples()
              data["samples_rec"] = samples_rec
              data["wav_rec"] = wav_rec
          if self.use_done_detector:
              samples_cut, wav_cut = self.cut_samples(samples, wav)
              data["samples_cut"] = samples_cut
              data["wav_cut"] = wav_cut
          self.log_epoch(data, epoch, stage)
  
imports: import logging
  import sys
  import torch
  import speechbrain
  import os
  from collections import namedtuple
  from audiomnist_prepare import prepare_audiomnist
  from hyperpyyaml import load_hyperpyyaml
  from torchaudio import functional
  from speechbrain.dataio.dataio import length_to_mask, write_audio
  from speechbrain.dataio.dataset import apply_overfit_test
  from speechbrain.utils import data_utils
  from speechbrain.utils.train_logger import plot_spectrogram
  from enum import Enum
  from speechbrain.utils.distributed import run_on_main
  from speechbrain.utils.data_utils import dict_value_combinations, dist_stats, masked_mean, masked_min, masked_max, masked_std, match_shape
label_desc: Use image node_id 36 for calling the DiffusionBrain obj's underlying member method code with example usage: obj.on_stage_end(stage, stage_loss, epoch) without return types