code: def training_step(self, batch, batch_idx):
      speakers = None
      if SpeakerID in self._train_dl.dataset.sup_data_types_set:
          (audio, audio_len, text, text_len, speakers) = batch
      else:
          (audio, audio_len, text, text_len) = batch
  
      spec, spec_lengths = self.audio_to_melspec_processor(
          audio, audio_len, linear_spec=True
      )
  
      with autocast(enabled=True):
          (
              audio_pred,
              l_length,
              attn,
              ids_slice,
              text_mask,
              z_mask,
              (z, z_p, m_p, logs_p, m_q, logs_q),
          ) = self.net_g(text, text_len, spec, spec_lengths, speakers)
  
      audio_pred = audio_pred.float()
  
      audio_pred_mel, _ = self.audio_to_melspec_processor(
          audio_pred.squeeze(1), audio_len, linear_spec=False
      )
  
      audio = slice_segments(
          audio.unsqueeze(1),
          ids_slice * self.cfg.n_window_stride,
          self._cfg.segment_size,
      )
      audio_mel, _ = self.audio_to_melspec_processor(
          audio.squeeze(1), audio_len, linear_spec=False
      )
  
      with autocast(enabled=True):
          y_d_hat_r, y_d_hat_g, _, _ = self.net_d(
              audio, audio_pred.detach()
          )
  
      with autocast(enabled=False):
          loss_disc, losses_disc_r, losses_disc_g = self.disc_loss(
              disc_real_outputs=y_d_hat_r,
              disc_generated_outputs=y_d_hat_g,
          )
          loss_disc_all = loss_disc
  
      # get optimizers
      optim_g, optim_d = self.optimizers()
  
      # train discriminator
      optim_d.zero_grad()
      self.manual_backward(loss_disc_all)
      norm_d = clip_grad_value_(self.net_d.parameters(), None)
      optim_d.step()
  
      with autocast(enabled=True):
          y_d_hat_r, y_d_hat_g, fmap_r, fmap_g = self.net_d(
              audio, audio_pred
          )
      # Generator
      with autocast(enabled=False):
          loss_dur = torch.sum(l_length.float())
          loss_mel = (
              F.l1_loss(audio_mel, audio_pred_mel) * self._cfg.c_mel
          )
          loss_kl = (
              self.kl_loss(
                  z_p=z_p,
                  logs_q=logs_q,
                  m_p=m_p,
                  logs_p=logs_p,
                  z_mask=z_mask,
              )
              * self._cfg.c_kl
          )
          loss_fm = self.feat_matching_loss(
              fmap_r=fmap_r, fmap_g=fmap_g
          )
          loss_gen, losses_gen = self.gen_loss(disc_outputs=y_d_hat_g)
          loss_gen_all = (
              loss_gen + loss_fm + loss_mel + loss_dur + loss_kl
          )
  
      # train generator
      optim_g.zero_grad()
      self.manual_backward(loss_gen_all)
      norm_g = clip_grad_value_(self.net_g.parameters(), None)
      optim_g.step()
  
      schedulers = self.lr_schedulers()
      if schedulers is not None:
          sch1, sch2 = schedulers
          if (
              self.trainer.is_last_batch
              and isinstance(
                  sch1, torch.optim.lr_scheduler.ExponentialLR
              )
              or isinstance(sch1, CosineAnnealing)
          ):
              sch1.step()
              sch2.step()
  
      metrics = {
          "loss_gen": loss_gen,
          "loss_fm": loss_fm,
          "loss_mel": loss_mel,
          "loss_dur": loss_dur,
          "loss_kl": loss_kl,
          "loss_gen_all": loss_gen_all,
          "loss_disc_all": loss_disc_all,
          "grad_gen": norm_g,
          "grad_disc": norm_d,
      }
  
      for i, v in enumerate(losses_gen):
          metrics[f"loss_gen_i_{i}"] = v
  
      for i, v in enumerate(losses_disc_r):
          metrics[f"loss_disc_r_{i}"] = v
  
      for i, v in enumerate(losses_disc_g):
          metrics[f"loss_disc_g_{i}"] = v
  
      self.log_dict(metrics, on_step=True, sync_dist=True)
  
imports: import contextlib
  import omegaconf
  import torch
  from hydra.utils import instantiate
  from omegaconf import DictConfig, OmegaConf
  from pytorch_lightning import Trainer
  from pytorch_lightning.loggers import WandbLogger
  from torch.cuda.amp import autocast
  from torch.nn import functional
  from nemo.collections.tts.data.dataset import DistributedBucketSampler
  from nemo.collections.tts.losses.vits_losses import DiscriminatorLoss, FeatureMatchingLoss, GeneratorLoss, KlLoss
  from nemo.collections.tts.models.base import TextToWaveform
  from nemo.collections.tts.modules.vits_modules import MultiPeriodDiscriminator
  from nemo.collections.tts.parts.utils.helpers import clip_grad_value_, g2p_backward_compatible_support, plot_spectrogram_to_numpy, slice_segments
  from nemo.collections.tts.torch.tts_data_types import SpeakerID
  from nemo.core.classes.common import PretrainedModelInfo, typecheck
  from nemo.core.neural_types.elements import AudioSignal, FloatType, Index, IntType, TokenIndex
  from nemo.core.neural_types.neural_type import NeuralType
  from nemo.core.optim.lr_scheduler import CosineAnnealing
  from nemo.utils import logging, model_utils
  from nemo.utils.decorators.experimental import experimental
label_desc: Use image node_id 7 for calling the VitsModel obj's underlying member method code with example usage: obj.training_step(batch, batch_idx) without return types