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: Saves end-of-epoch logs
  
  Arguments
  ---------
  data: dict
      the data to be loged, with the following keys
      samples: generated samples
      wav: generated waveform
      samples_rec: reconstruction samples (to assess autoencoder quality)
      samples_wav: reconstruction audio (to assess autoencoder quality)
  epoch: int
      the epoch number
  stage: speechbrain.Stage
      the training stage
docstr_tok: ["Saves","end-of-epoch","logs","Arguments","--","--","--","--","-","data",":","dict","the","data","to","be","loged",",","with","the","following","keys","samples",":","generated","samples","wav",":","generated","waveform","samples_rec",":","reconstruction","samples","(","to","assess","autoencoder","quality",")","samples_wav",":","reconstruction","audio","(","to","assess","autoencoder","quality",")","epoch",":","int","the","epoch","number","stage",":","speechbrain.Stage","the","training","stage"]
code: def log_epoch(self, data, epoch, stage):
      """Saves end-of-epoch logs
  
      Arguments
      ---------
      data: dict
          the data to be loged, with the following keys
          samples: generated samples
          wav: generated waveform
          samples_rec: reconstruction samples (to assess autoencoder quality)
          samples_wav: reconstruction audio (to assess autoencoder quality)
      epoch: int
          the epoch number
      stage: speechbrain.Stage
          the training stage
  
      """
      epoch_sample_path = os.path.join(
          self.hparams.sample_folder, str(epoch)
      )
      samples, samples_denorm, wav, labels, samples_rec, wav_rec = [
          data.get(key)
          for key in [
              "samples",
              "samples_denorm",
              "wav",
              "labels",
              "samples_rec",
              "wav_rec",
          ]
      ]
      if not torch.is_tensor(samples):
          samples = torch.stack(samples)
      samples_log = data.get("samples_cut", samples)
      wav_log = data.get("wav_cut", wav)
      self.save_spectrograms(
          samples_log, epoch_sample_path, labels=labels
      )
      sample_ids = torch.arange(1, len(samples) + 1)
      for metric in self.sample_metrics:
          metric.append(sample_ids, samples)
      if wav is not None:
          self.save_audio(wav_log, epoch_sample_path, labels=labels)
      if self.diffusion_mode == DiffusionMode.LATENT:
          self.save_spectrograms(
              samples_rec, epoch_sample_path, folder="spec_rec"
          )
          if wav_rec is not None:
              self.save_audio(
                  wav_rec, epoch_sample_path, folder="wav_rec"
              )
  
      self.save_raw(
          spec=samples,
          spec_denorm=samples_denorm,
          wav=wav,
          spec_rec=samples_rec,
          wav_rec=wav_rec,
          path=epoch_sample_path,
      )
      if self.hparams.use_tensorboard:
          sample_mean_stats = self.sample_mean_metric.summarize()
          sample_std_stats = self.sample_std_metric.summarize()
          sample_min_stats = self.sample_min_metric.summarize()
          sample_max_stats = self.sample_max_metric.summarize()
          stats = {
              "sample_mean": sample_mean_stats["average"],
              "sample_mean_min": sample_mean_stats["min_score"],
              "sample_mean_max": sample_mean_stats["max_score"],
              "sample_std": sample_std_stats["average"],
              "sample_std_min": sample_std_stats["min_score"],
              "sample_std_max": sample_std_stats["max_score"],
              "sample_min": sample_min_stats["min_score"],
              "sample_max": sample_max_stats["max_score"],
          }
          stats_args = {f"{stage.name.lower()}_stats": stats}
          self.hparams.tensorboard_train_logger.log_stats(
              stats_meta={"step": self.step}, **stats_args
          )
          self.log_samples(
              spectrogram_samples=samples_log, wav_samples=wav_log
          )
          if self.diffusion_mode == DiffusionMode.LATENT:
              self.log_samples(
                  spectrogram_samples=samples_rec,
                  wav_samples=wav_rec,
                  key_prefix="rec_",
              )
  
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 38 for calling the DiffusionBrain obj's underlying member method code with example usage: obj.log_epoch(data, epoch, stage) without return types