docstr: Compute gradients and update either D or G based on sub-stage.
docstr_tok: ["Compute","gradients","and","update","either","D","or","G","based","on","sub-stage","."]
code: def fit_batch(self, batch):
      "Compute gradients and update either D or G based on sub-stage."
      predictions = self.compute_forward(batch, sb.Stage.TRAIN)
      loss_tracker = 0
      if self.sub_stage == SubStage.CURRENT:
          for mode in ["clean", "enh", "noisy"]:
              loss = self.compute_objectives(
                  predictions, batch, sb.Stage.TRAIN, f"D_{mode}"
              )
              self.d_optimizer.zero_grad()
              loss.backward()
              torch.nn.utils.clip_grad_norm_(
                  self.modules.parameters(), self.max_grad_norm
              )
              self.d_optimizer.step()
              loss_tracker += loss.detach() / 3
      elif self.sub_stage == SubStage.HISTORICAL:
          loss = self.compute_objectives(
              predictions, batch, sb.Stage.TRAIN, "D_enh"
          )
          self.d_optimizer.zero_grad()
          loss.backward()
          torch.nn.utils.clip_grad_norm_(
              self.modules.parameters(), self.max_grad_norm
          )
          self.d_optimizer.step()
          loss_tracker += loss.detach()
      elif self.sub_stage == SubStage.GENERATOR:
          for name, param in self.modules.generator.named_parameters():
              if "Learnable_sigmoid" in name:
                  param.data = torch.clamp(
                      param, max=3.5
                  )  # to prevent gradient goes to infinity
                  param.data[param != param] = 3.5  # set 'nan' to 3.5
  
          loss = self.compute_objectives(
              predictions, batch, sb.Stage.TRAIN, "generator"
          )
  
          self.g_optimizer.zero_grad()
          loss.backward()
          torch.nn.utils.clip_grad_norm_(
              self.modules.parameters(), self.max_grad_norm
          )
          self.g_optimizer.step()
          loss_tracker += loss.detach()
  
      return loss_tracker
  
imports: import os
  import sys
  import shutil
  import pickle
  import torch
  import torchaudio
  import speechbrain
  from pesq import pesq
  from enum import Enum, auto
  from hyperpyyaml import load_hyperpyyaml
  from speechbrain.utils.metric_stats import MetricStats
  from speechbrain.processing.features import spectral_magnitude
  from speechbrain.nnet.loss.stoi_loss import stoi_loss
  from speechbrain.utils.distributed import run_on_main
  from speechbrain.dataio.sampler import ReproducibleWeightedRandomSampler
label_desc: Use image node_id 10 for calling the MetricGanBrain obj's underlying member method code with example usage: obj.fit_batch(batch) and returns: loss_tracker