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
  
      if self.hparams.use_whamr_train:
          whamr_prob = torch.rand(1).item()
          if whamr_prob > (1 - self.hparams.whamr_proportion):
              batch = next(self.hparams.train_whamr_loader)
  
      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,
                      snrhat,
                      snr,
                      snr_compressed,
                  ) = self.compute_forward(
                      mixture, targets, sb.Stage.TRAIN, noise
                  )
  
                  snr = snr.reshape(-1)
                  loss = ((snr_compressed - snrhat).abs()).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:
              # get the oracle snrs, estimated snrs, and the source estimates
              (
                  predictions,
                  snrhat,
                  snr,
                  snr_compressed,
              ) = self.compute_forward(
                  mixture, targets, sb.Stage.TRAIN, noise
              )
  
              snr = snr.reshape(-1)
              loss = ((snr_compressed - snrhat).abs()).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 speechbrain
  import speechbrain.nnet.schedulers
  from speechbrain.utils.distributed import run_on_main
  from hyperpyyaml import load_hyperpyyaml
  import itertools
  from tqdm import tqdm
  import numpy
  import logging
  import csv
  from speechbrain.core import AMPConfig
label_desc: Use image node_id 5 for calling the Separation obj's underlying member method code with example usage: obj.fit_batch(batch) and returns: loss