code: def main(args):  # pylint: disable=redefined-outer-name
      # pylint: disable=global-variable-undefined
      global train_data, eval_data
  
      # setup audio processor
      ap = AudioProcessor(**c.audio)
  
      # 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:
      #     mel_feat_path = os.path.join(OUT_PATH, "mel")
      #     feat_data = find_feat_files(mel_feat_path)
      #     if feat_data:
      #         print(f" > Loading features from: {mel_feat_path}")
      #         eval_data, train_data = load_wav_feat_data(
      #             c.data_path, mel_feat_path, c.eval_split_size
      #         )
      #     else:
      #         print(" > No feature data found. Preprocessing...")
      #         # preprocessing feature data from given wav files
      #         preprocess_wav_files(OUT_PATH, CONFIG, ap)
      #         eval_data, train_data = load_wav_feat_data(
      #             c.data_path, mel_feat_path, c.eval_split_size
      #         )
  
      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 model
      model_wavernn = setup_wavernn(c)
  
      # setup amp scaler
      scaler = (
          torch.cuda.amp.GradScaler() if c.mixed_precision else None
      )
  
      # define train functions
      if c.mode == "mold":
          criterion = discretized_mix_logistic_loss
      elif c.mode == "gauss":
          criterion = gaussian_loss
      elif isinstance(c.mode, int):
          criterion = torch.nn.CrossEntropyLoss()
  
      if use_cuda:
          model_wavernn.cuda()
          if isinstance(c.mode, int):
              criterion.cuda()
  
      optimizer = RAdam(
          model_wavernn.parameters(), lr=c.lr, weight_decay=0
      )
  
      scheduler = None
      if "lr_scheduler" in c:
          scheduler = getattr(torch.optim.lr_scheduler, c.lr_scheduler)
          scheduler = scheduler(optimizer, **c.lr_scheduler_params)
      # slow start for the first 5 epochs
      # lr_lambda = lambda epoch: min(epoch / c.warmup_steps, 1)
      # scheduler = optim.lr_scheduler.LambdaLR(optimizer, lr_lambda)
  
      # restore any checkpoint
      if args.restore_path:
          checkpoint = torch.load(args.restore_path, map_location="cpu")
          try:
              print(" > Restoring Model...")
              model_wavernn.load_state_dict(checkpoint["model"])
              print(" > Restoring Optimizer...")
              optimizer.load_state_dict(checkpoint["optimizer"])
              if "scheduler" in checkpoint:
                  print(" > Restoring Generator LR Scheduler...")
                  scheduler.load_state_dict(checkpoint["scheduler"])
                  scheduler.optimizer = optimizer
              if "scaler" in checkpoint and c.mixed_precision:
                  print(" > Restoring AMP Scaler...")
                  scaler.load_state_dict(checkpoint["scaler"])
          except RuntimeError:
              # retore only matching layers.
              print(" > Partial model initialization...")
              model_dict = model_wavernn.state_dict()
              model_dict = set_init_dict(
                  model_dict, checkpoint["model"], c
              )
              model_wavernn.load_state_dict(model_dict)
  
          print(
              " > Model restored from step %d" % checkpoint["step"],
              flush=True,
          )
          args.restore_step = checkpoint["step"]
      else:
          args.restore_step = 0
  
      # DISTRIBUTED
      # if num_gpus > 1:
      #     model = apply_gradient_allreduce(model)
  
      num_parameters = count_parameters(model_wavernn)
      print(
          " > Model has {} parameters".format(num_parameters),
          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_wavernn,
              optimizer,
              criterion,
              scheduler,
              scaler,
              ap,
              global_step,
              epoch,
          )
          eval_avg_loss_dict = evaluate(
              model_wavernn, criterion, ap, global_step, epoch
          )
          c_logger.print_epoch_end(epoch, eval_avg_loss_dict)
          target_loss = eval_avg_loss_dict["avg_model_loss"]
          best_loss = save_best_model(
              target_loss,
              best_loss,
              model_wavernn,
              optimizer,
              scheduler,
              None,
              None,
              None,
              global_step,
              epoch,
              OUT_PATH,
              model_losses=eval_avg_loss_dict,
              scaler=scaler.state_dict() if c.mixed_precision else None,
          )
  
imports: import argparse
  import os
  import sys
  import traceback
  import time
  import glob
  import random
  import torch
  from torch.utils.data import DataLoader
  from TTS.tts.utils.visual import plot_spectrogram
  from TTS.utils.audio import AudioProcessor
  from TTS.utils.radam import RAdam
  from TTS.utils.io import copy_model_files, load_config
  from TTS.utils.training import setup_torch_training_env
  from TTS.utils.console_logger import ConsoleLogger
  from TTS.utils.tensorboard_logger import TensorboardLogger
  from TTS.utils.generic_utils import KeepAverage, count_parameters, create_experiment_folder, get_git_branch, remove_experiment_folder, set_init_dict
  from TTS.vocoder.datasets.wavernn_dataset import WaveRNNDataset
  from TTS.vocoder.datasets.preprocess import load_wav_data, load_wav_feat_data
  from TTS.vocoder.utils.distribution import discretized_mix_logistic_loss, gaussian_loss
  from TTS.vocoder.utils.generic_utils import setup_wavernn
  from TTS.vocoder.utils.io import save_best_model, save_checkpoint
label_desc: Use image node_id 5 for calling a global function with example usage: main(args) without return types