docstr: Trains one batch
docstr_tok: ["Trains","one","batch"]
code: def fit_batch(self, batch):
      """Trains one batch"""
      amp = AMPConfig.from_name(self.precision)
      should_step = (self.step % self.grad_accumulation_factor) == 0
  
      # Unpacking batch list
      mixture = batch.mix_sig
      targets = [batch.s1_sig, batch.s2_sig]
      if self.hparams.use_wham_noise:
          noise = batch.noise_sig[0]
      else:
          noise = None
  
      if self.hparams.num_spks == 3:
          targets.append(batch.s3_sig)
  
      with self.no_sync(not should_step):
          if self.use_amp:
              with torch.autocast(
                  dtype=amp.dtype,
                  device_type=torch.device(self.device).type,
              ):
                  predictions, targets = self.compute_forward(
                      mixture, targets, sb.Stage.TRAIN, noise
                  )
                  loss = self.compute_objectives(predictions, targets)
  
                  # hard threshold the easy dataitems
                  if self.hparams.threshold_byloss:
                      th = self.hparams.threshold
                      loss = loss[loss > th]
                      if loss.nelement() > 0:
                          loss = loss.mean()
                  else:
                      loss = loss.mean()
  
              if (
                  loss.nelement() > 0
                  and loss < self.hparams.loss_upper_lim
              ):  # the fix for computational problems
                  self.scaler.scale(loss).backward()
                  if self.hparams.clip_grad_norm >= 0:
                      self.scaler.unscale_(self.optimizer)
                      torch.nn.utils.clip_grad_norm_(
                          self.modules.parameters(),
                          self.hparams.clip_grad_norm,
                      )
                  self.scaler.step(self.optimizer)
                  self.scaler.update()
              else:
                  self.nonfinite_count += 1
                  logger.info(
                      "infinite loss or empty loss! it happened {} times so far - skipping this batch".format(
                          self.nonfinite_count
                      )
                  )
                  loss.data = torch.tensor(0).to(self.device)
          else:
              predictions, targets = self.compute_forward(
                  mixture, targets, sb.Stage.TRAIN, noise
              )
              loss = self.compute_objectives(predictions, targets)
  
              if self.hparams.threshold_byloss:
                  th = self.hparams.threshold
                  loss = loss[loss > th]
                  if loss.nelement() > 0:
                      loss = loss.mean()
              else:
                  loss = loss.mean()
  
              if (
                  loss.nelement() > 0
                  and loss < self.hparams.loss_upper_lim
              ):  # the fix for computational problems
                  loss.backward()
                  if self.hparams.clip_grad_norm >= 0:
                      torch.nn.utils.clip_grad_norm_(
                          self.modules.parameters(),
                          self.hparams.clip_grad_norm,
                      )
                  self.optimizer.step()
              else:
                  self.nonfinite_count += 1
                  logger.info(
                      "infinite loss or empty loss! it happened {} times so far - skipping this batch".format(
                          self.nonfinite_count
                      )
                  )
                  loss.data = torch.tensor(0).to(self.device)
      self.optimizer.zero_grad()
  
      return loss.detach().cpu()
  
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 3 for calling the Separation obj's underlying member method code with example usage: obj.fit_batch(batch) and returns: loss