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: Train the parameters given a single batch in input
docstr_tok: ["Train","the","parameters","given","a","single","batch","in","input"]
code: def fit_batch(self, batch):
      """Train the parameters given a single batch in input"""
      if self.reference_batch is None:
          self.reference_batch = batch
  
      should_step = self.step % self.grad_accumulation_factor == 0
      outputs = self.compute_forward(batch, sb.Stage.TRAIN)
      loss, loss_autoencoder, loss_done = self.compute_objectives(
          outputs, batch, sb.Stage.TRAIN
      )
      if self.train_diffusion:
          with self.no_sync(not should_step):
              (loss / self.grad_accumulation_factor).backward(
                  retain_graph=True
              )
      # Done loss - iffapplicable
      if self.use_done_detector:
          with self.no_sync(not should_step):
              (loss_done / self.grad_accumulation_factor).backward(
                  retain_graph=True
              )
  
      if should_step:
          if self.train_diffusion:
              self.optimizer.step()
          self.optimizer.zero_grad()
  
          if self.use_done_detector:
              self.optimizer_done.step()
              self.optimizer_done.zero_grad()
  
      # Latent diffusion: Step through the autoencoder
      if (
          self.diffusion_mode == DiffusionMode.LATENT
          and loss_autoencoder is not None
      ):
          with self.no_sync(not should_step):
              (
                  loss_autoencoder / self.grad_accumulation_factor
              ).backward()
          if should_step:
              torch.nn.utils.clip_grad_norm_(
                  self.modules.parameters(), self.max_grad_norm
              )
              self.autoencoder_optimizer.step()
          self.autoencoder_optimizer.zero_grad()
  
      self.optimizer_step += 1
      self.hparams.lr_annealing(self.optimizer, self.optimizer_step)
      if self.diffusion_mode == DiffusionMode.LATENT:
          self.hparams.lr_annealing_autoencoder(
              self.autoencoder_optimizer, self.optimizer_step
          )
      if (
          self.hparams.enable_train_metrics
          and self.hparams.use_tensorboard
          and (
              self.step == 1
              or self.step % self.hparams.train_log_interval == 0
          )
      ):
          self.log_batch(outputs)
      return loss
  
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 8 for calling the DiffusionBrain obj's underlying member method code with example usage: obj.fit_batch(batch) and returns: loss