class_docstr: UnivNet model (https://arxiv.org/abs/2106.07889) that is used to generate audio from mel spectrogram.
class_docstr_tok: ["UnivNet","model","(","https",":","\/\/arxiv.org\/abs\/2106.07889",")","that","is","used","to","generate","audio","from","mel","spectrogram","."]
code: def training_step(self, batch, batch_idx):
      if self.input_as_mel:
          # Pre-computed spectrograms will be used as input
          audio, audio_len, audio_mel = batch
      else:
          audio, audio_len = batch
          audio_mel, _ = self.audio_to_melspec_precessor(
              audio, audio_len
          )
  
      audio = audio.unsqueeze(1)
  
      audio_pred = self.generator(x=audio_mel)
      audio_pred_mel, _ = self.trg_melspec_fn(
          audio_pred.squeeze(1), audio_len
      )
  
      optim_g, optim_d = self.optimizers()
  
      # Train discriminator
      optim_d.zero_grad()
      mpd_score_real, mpd_score_gen, _, _ = self.mpd(
          y=audio, y_hat=audio_pred.detach()
      )
      loss_disc_mpd, _, _ = self.discriminator_loss(
          disc_real_outputs=mpd_score_real,
          disc_generated_outputs=mpd_score_gen,
      )
      mrd_score_real, mrd_score_gen, _, _ = self.mrd(
          y=audio, y_hat=audio_pred.detach()
      )
      loss_disc_mrd, _, _ = self.discriminator_loss(
          disc_real_outputs=mrd_score_real,
          disc_generated_outputs=mrd_score_gen,
      )
      loss_d = loss_disc_mrd + loss_disc_mpd
      self.manual_backward(loss_d)
      optim_d.step()
  
      # Train generator
      optim_g.zero_grad()
      loss_sc, loss_mag = self.mrstft_loss(
          x=audio_pred.squeeze(1),
          y=audio.squeeze(1),
          input_lengths=audio_len,
      )
      loss_sc = torch.stack(loss_sc).mean()
      loss_mag = torch.stack(loss_mag).mean()
      loss_mrstft = (loss_sc + loss_mag) * self.stft_lamb
      _, mpd_score_gen, _, _ = self.mpd(y=audio, y_hat=audio_pred)
      _, mrd_score_gen, _, _ = self.mrd(y=audio, y_hat=audio_pred)
      loss_gen_mpd, _ = self.generator_loss(disc_outputs=mpd_score_gen)
      loss_gen_mrd, _ = self.generator_loss(disc_outputs=mrd_score_gen)
      loss_g = loss_gen_mrd + loss_gen_mpd + loss_mrstft
      self.manual_backward(loss_g)
      optim_g.step()
  
      metrics = {
          "g_loss_sc": loss_sc,
          "g_loss_mag": loss_mag,
          "g_loss_mrstft": loss_mrstft,
          "g_loss_gen_mpd": loss_gen_mpd,
          "g_loss_gen_mrd": loss_gen_mrd,
          "g_loss": loss_g,
          "d_loss_mpd": loss_disc_mpd,
          "d_loss_mrd": loss_disc_mrd,
          "d_loss": loss_d,
          "global_step": self.global_step,
          "lr": optim_g.param_groups[0]["lr"],
      }
      self.log_dict(metrics, on_step=True, sync_dist=True)
      self.log(
          "g_mrstft_loss",
          loss_mrstft,
          prog_bar=True,
          logger=False,
          sync_dist=True,
      )
  
imports: import itertools
  from typing import Dict
  import torch
  import torch.nn.functional
  from hydra.utils import instantiate
  from omegaconf import DictConfig, OmegaConf, open_dict
  from pytorch_lightning.loggers.wandb import WandbLogger
  from nemo.collections.tts.losses.hifigan_losses import DiscriminatorLoss, GeneratorLoss
  from nemo.collections.tts.losses.stftlosses import MultiResolutionSTFTLoss
  from nemo.collections.tts.models.base import Vocoder
  from nemo.collections.tts.modules.univnet_modules import MultiPeriodDiscriminator, MultiResolutionDiscriminator
  from nemo.collections.tts.parts.utils.helpers import get_batch_size, get_num_workers, plot_spectrogram_to_numpy
  from nemo.core import Exportable
  from nemo.core.classes.common import PretrainedModelInfo, typecheck
  from nemo.core.neural_types.elements import AudioSignal, MelSpectrogramType
  from nemo.core.neural_types.neural_type import NeuralType
  from nemo.core.optim.lr_scheduler import compute_max_steps, prepare_lr_scheduler
  from nemo.utils import logging, model_utils
label_desc: Use image node_id 7 for calling the UnivNetModel obj's underlying member method code with example usage: obj.training_step(batch, batch_idx) without return types