docstr: Compute model outputs and the loss values. `optimizer_idx` selects the generator or the discriminator for
  network on the current pass.
  
  Args:
      batch (Dict): Batch of samples returned by the dataloader.
      criterion (Dict): Criterion used to compute the losses.
      optimizer_idx (int): ID of the optimizer in use on the current pass.
  
  Raises:
      ValueError: `optimizer_idx` is an unexpected value.
  
  Returns:
      Tuple[Dict, Dict]: model outputs and the computed loss values.
docstr_tok: ["Compute","model","outputs","and","the","loss","values",".","`","optimizer_idx","`","selects","the","generator","or","the","discriminator","for","network","on","the","current","pass",".","Args",":","batch","(","Dict",")",":","Batch","of","samples","returned","by","the","dataloader",".","criterion","(","Dict",")",":","Criterion","used","to","compute","the","losses",".","optimizer_idx","(","int",")",":","ID","of","the","optimizer","in","use","on","the","current","pass",".","Raises",":","ValueError",":","`","optimizer_idx","`","is","an","unexpected","value",".","Returns",":","Tuple","[","Dict",",","Dict","]",":","model","outputs","and","the","computed","loss","values","."]
code: def train_step(
      self, batch: Dict, criterion: Dict, optimizer_idx: int
  ) -> Tuple[Dict, Dict]:
      """Compute model outputs and the loss values. `optimizer_idx` selects the generator or the discriminator for
      network on the current pass.
  
      Args:
          batch (Dict): Batch of samples returned by the dataloader.
          criterion (Dict): Criterion used to compute the losses.
          optimizer_idx (int): ID of the optimizer in use on the current pass.
  
      Raises:
          ValueError: `optimizer_idx` is an unexpected value.
  
      Returns:
          Tuple[Dict, Dict]: model outputs and the computed loss values.
      """
      outputs = {}
      loss_dict = {}
  
      x = batch["input"]
      y = batch["waveform"]
  
      if optimizer_idx not in [0, 1]:
          raise ValueError(" [!] Unexpected `optimizer_idx`.")
  
      if optimizer_idx == 0:
          # DISCRIMINATOR optimization
  
          # generator pass
          y_hat = self.model_g(x)[:, :, : y.size(2)]
  
          # cache for generator loss
          # pylint: disable=W0201
          self.y_hat_g = y_hat
          self.y_hat_sub = None
          self.y_sub_g = None
  
          # PQMF formatting
          if y_hat.shape[1] > 1:
              self.y_hat_sub = y_hat
              y_hat = self.model_g.pqmf_synthesis(y_hat)
              self.y_hat_g = y_hat  # save for generator loss
              self.y_sub_g = self.model_g.pqmf_analysis(y)
  
          scores_fake, feats_fake, feats_real = None, None, None
  
          if self.train_disc:
              # use different samples for G and D trainings
              if self.config.diff_samples_for_G_and_D:
                  x_d = batch["input_disc"]
                  y_d = batch["waveform_disc"]
                  # use a different sample than generator
                  with torch.no_grad():
                      y_hat = self.model_g(x_d)
  
                  # PQMF formatting
                  if y_hat.shape[1] > 1:
                      y_hat = self.model_g.pqmf_synthesis(y_hat)
              else:
                  # use the same samples as generator
                  x_d = x.clone()
                  y_d = y.clone()
                  y_hat = self.y_hat_g
  
              # run D with or without cond. features
              if len(signature(self.model_d.forward).parameters) == 2:
                  D_out_fake = self.model_d(y_hat.detach().clone(), x_d)
                  D_out_real = self.model_d(y_d, x_d)
              else:
                  D_out_fake = self.model_d(y_hat.detach())
                  D_out_real = self.model_d(y_d)
  
              # format D outputs
              if isinstance(D_out_fake, tuple):
                  # self.model_d returns scores and features
                  scores_fake, feats_fake = D_out_fake
                  if D_out_real is None:
                      scores_real, feats_real = None, None
                  else:
                      scores_real, feats_real = D_out_real
              else:
                  # model D returns only scores
                  scores_fake = D_out_fake
                  scores_real = D_out_real
  
              # compute losses
              loss_dict = criterion[optimizer_idx](
                  scores_fake, scores_real
              )
              outputs = {"model_outputs": y_hat}
  
      if optimizer_idx == 1:
          # GENERATOR loss
          scores_fake, feats_fake, feats_real = None, None, None
          if self.train_disc:
              if len(signature(self.model_d.forward).parameters) == 2:
                  D_out_fake = self.model_d(self.y_hat_g, x)
              else:
                  D_out_fake = self.model_d(self.y_hat_g)
              D_out_real = None
  
              if self.config.use_feat_match_loss:
                  with torch.no_grad():
                      D_out_real = self.model_d(y)
  
              # format D outputs
              if isinstance(D_out_fake, tuple):
                  scores_fake, feats_fake = D_out_fake
                  if D_out_real is None:
                      feats_real = None
                  else:
                      _, feats_real = D_out_real
              else:
                  scores_fake = D_out_fake
                  feats_fake, feats_real = None, None
  
          # compute losses
          loss_dict = criterion[optimizer_idx](
              self.y_hat_g,
              y,
              scores_fake,
              feats_fake,
              feats_real,
              self.y_hat_sub,
              self.y_sub_g,
          )
          outputs = {"model_outputs": self.y_hat_g}
      return outputs, loss_dict
  
imports: from inspect import signature
  from typing import Dict, List, Tuple
  import numpy
  import torch
  from coqpit import Coqpit
  from torch import nn
  from torch.utils.data import DataLoader
  from torch.utils.data.distributed import DistributedSampler
  from trainer.trainer_utils import get_optimizer, get_scheduler
  from TTS.utils.audio import AudioProcessor
  from TTS.utils.io import load_fsspec
  from TTS.vocoder.datasets.gan_dataset import GANDataset
  from TTS.vocoder.layers.losses import DiscriminatorLoss, GeneratorLoss
  from TTS.vocoder.models import setup_discriminator, setup_generator
  from TTS.vocoder.models.base_vocoder import BaseVocoder
  from TTS.vocoder.utils.generic_utils import plot_results
label_desc: Use image node_id 4 for calling the GAN obj's underlying member method code with example usage: obj.train_step(batch, criterion, optimizer_idx) and returns: outputs,  loss_dict