code: def main(args):  # pylint: disable=redefined-outer-name
      # pylint: disable=global-variable-undefined
      global train_data, eval_data
      print(f" > Loading wavs from: {c.data_path}")
      if c.feature_path is not None:
          print(f" > Loading features from: {c.feature_path}")
          eval_data, train_data = load_wav_feat_data(
              c.data_path, c.feature_path, c.eval_split_size
          )
      else:
          eval_data, train_data = load_wav_data(
              c.data_path, c.eval_split_size
          )
  
      # setup audio processor
      ap = AudioProcessor(**c.audio)
  
      # DISTRUBUTED
      if num_gpus > 1:
          init_distributed(
              args.rank,
              num_gpus,
              args.group_id,
              c.distributed["backend"],
              c.distributed["url"],
          )
  
      # setup models
      model_gen = setup_generator(c)
      model_disc = setup_discriminator(c)
  
      # setup optimizers
      optimizer_gen = RAdam(
          model_gen.parameters(), lr=c.lr_gen, weight_decay=0
      )
      optimizer_disc = RAdam(
          model_disc.parameters(), lr=c.lr_disc, weight_decay=0
      )
  
      # schedulers
      scheduler_gen = None
      scheduler_disc = None
      if "lr_scheduler_gen" in c:
          scheduler_gen = getattr(
              torch.optim.lr_scheduler, c.lr_scheduler_gen
          )
          scheduler_gen = scheduler_gen(
              optimizer_gen, **c.lr_scheduler_gen_params
          )
      if "lr_scheduler_disc" in c:
          scheduler_disc = getattr(
              torch.optim.lr_scheduler, c.lr_scheduler_disc
          )
          scheduler_disc = scheduler_disc(
              optimizer_disc, **c.lr_scheduler_disc_params
          )
  
      # setup criterion
      criterion_gen = GeneratorLoss(c)
      criterion_disc = DiscriminatorLoss(c)
  
      if args.restore_path:
          checkpoint = torch.load(args.restore_path, map_location="cpu")
          try:
              print(" > Restoring Generator Model...")
              model_gen.load_state_dict(checkpoint["model"])
              print(" > Restoring Generator Optimizer...")
              optimizer_gen.load_state_dict(checkpoint["optimizer"])
              print(" > Restoring Discriminator Model...")
              model_disc.load_state_dict(checkpoint["model_disc"])
              print(" > Restoring Discriminator Optimizer...")
              optimizer_disc.load_state_dict(
                  checkpoint["optimizer_disc"]
              )
              if "scheduler" in checkpoint:
                  print(" > Restoring Generator LR Scheduler...")
                  scheduler_gen.load_state_dict(checkpoint["scheduler"])
                  # NOTE: Not sure if necessary
                  scheduler_gen.optimizer = optimizer_gen
              if "scheduler_disc" in checkpoint:
                  print(" > Restoring Discriminator LR Scheduler...")
                  scheduler_disc.load_state_dict(
                      checkpoint["scheduler_disc"]
                  )
                  scheduler_disc.optimizer = optimizer_disc
          except RuntimeError:
              # retore only matching layers.
              print(" > Partial model initialization...")
              model_dict = model_gen.state_dict()
              model_dict = set_init_dict(
                  model_dict, checkpoint["model"], c
              )
              model_gen.load_state_dict(model_dict)
  
              model_dict = model_disc.state_dict()
              model_dict = set_init_dict(
                  model_dict, checkpoint["model_disc"], c
              )
              model_disc.load_state_dict(model_dict)
              del model_dict
  
          # reset lr if not countinuining training.
          for group in optimizer_gen.param_groups:
              group["lr"] = c.lr_gen
  
          for group in optimizer_disc.param_groups:
              group["lr"] = c.lr_disc
  
          print(
              " > Model restored from step %d" % checkpoint["step"],
              flush=True,
          )
          args.restore_step = checkpoint["step"]
      else:
          args.restore_step = 0
  
      if use_cuda:
          model_gen.cuda()
          criterion_gen.cuda()
          model_disc.cuda()
          criterion_disc.cuda()
  
      # DISTRUBUTED
      if num_gpus > 1:
          model_gen = DDP_th(model_gen, device_ids=[args.rank])
          model_disc = DDP_th(model_disc, device_ids=[args.rank])
  
      num_params = count_parameters(model_gen)
      print(
          " > Generator has {} parameters".format(num_params),
          flush=True,
      )
      num_params = count_parameters(model_disc)
      print(
          " > Discriminator has {} parameters".format(num_params),
          flush=True,
      )
  
      if "best_loss" not in locals():
          best_loss = float("inf")
  
      global_step = args.restore_step
      for epoch in range(0, c.epochs):
          c_logger.print_epoch_start(epoch, c.epochs)
          _, global_step = train(
              model_gen,
              criterion_gen,
              optimizer_gen,
              model_disc,
              criterion_disc,
              optimizer_disc,
              scheduler_gen,
              scheduler_disc,
              ap,
              global_step,
              epoch,
          )
          eval_avg_loss_dict = evaluate(
              model_gen,
              criterion_gen,
              model_disc,
              criterion_disc,
              ap,
              global_step,
              epoch,
          )
          c_logger.print_epoch_end(epoch, eval_avg_loss_dict)
          target_loss = eval_avg_loss_dict[c.target_loss]
          best_loss = save_best_model(
              target_loss,
              best_loss,
              model_gen,
              optimizer_gen,
              scheduler_gen,
              model_disc,
              optimizer_disc,
              scheduler_disc,
              global_step,
              epoch,
              OUT_PATH,
              model_losses=eval_avg_loss_dict,
          )
  
imports: import argparse
  import glob
  import os
  import sys
  import time
  import traceback
  from inspect import signature
  import torch
  from torch.utils.data import DataLoader
  from TTS.utils.audio import AudioProcessor
  from TTS.utils.console_logger import ConsoleLogger
  from TTS.utils.generic_utils import KeepAverage, count_parameters, create_experiment_folder, get_git_branch, remove_experiment_folder, set_init_dict
  from TTS.utils.io import copy_model_files, load_config
  from TTS.utils.radam import RAdam
  from TTS.utils.tensorboard_logger import TensorboardLogger
  from TTS.utils.training import setup_torch_training_env
  from TTS.vocoder.datasets.gan_dataset import GANDataset
  from TTS.vocoder.datasets.preprocess import load_wav_data, load_wav_feat_data
  from TTS.vocoder.layers.losses import DiscriminatorLoss, GeneratorLoss
  from TTS.vocoder.utils.generic_utils import plot_results, setup_discriminator, setup_generator
  from TTS.vocoder.utils.io import save_best_model, save_checkpoint
  from torch.nn.parallel import DistributedDataParallel
  from torch.utils.data.distributed import DistributedSampler
  from TTS.utils.distribute import init_distributed
label_desc: Use image node_id 5 for calling a global function with example usage: main(args) without return types