code: def train(
      model,
      optimizer,
      criterion,
      scheduler,
      scaler,
      ap,
      global_step,
      epoch,
  ):
      # create train loader
      data_loader = setup_loader(ap, is_val=False, verbose=(epoch == 0))
      model.train()
      epoch_time = 0
      keep_avg = KeepAverage()
      if use_cuda:
          batch_n_iter = int(
              len(data_loader.dataset) / (c.batch_size * num_gpus)
          )
      else:
          batch_n_iter = int(len(data_loader.dataset) / c.batch_size)
      end_time = time.time()
      c_logger.print_train_start()
      # train loop
      for num_iter, data in enumerate(data_loader):
          start_time = time.time()
          x_input, mels, y_coarse = format_data(data)
          loader_time = time.time() - end_time
          global_step += 1
  
          optimizer.zero_grad()
  
          if c.mixed_precision:
              # mixed precision training
              with torch.cuda.amp.autocast():
                  y_hat = model(x_input, mels)
                  if isinstance(model.mode, int):
                      y_hat = y_hat.transpose(1, 2).unsqueeze(-1)
                  else:
                      y_coarse = y_coarse.float()
                  y_coarse = y_coarse.unsqueeze(-1)
                  # compute losses
                  loss = criterion(y_hat, y_coarse)
              scaler.scale(loss).backward()
              scaler.unscale_(optimizer)
              if c.grad_clip > 0:
                  torch.nn.utils.clip_grad_norm_(
                      model.parameters(), c.grad_clip
                  )
              scaler.step(optimizer)
              scaler.update()
          else:
              # full precision training
              y_hat = model(x_input, mels)
              if isinstance(model.mode, int):
                  y_hat = y_hat.transpose(1, 2).unsqueeze(-1)
              else:
                  y_coarse = y_coarse.float()
              y_coarse = y_coarse.unsqueeze(-1)
              # compute losses
              loss = criterion(y_hat, y_coarse)
              if loss.item() is None:
                  raise RuntimeError(" [!] None loss. Exiting ...")
              loss.backward()
              if c.grad_clip > 0:
                  torch.nn.utils.clip_grad_norm_(
                      model.parameters(), c.grad_clip
                  )
              optimizer.step()
  
          if scheduler is not None:
              scheduler.step()
  
          # get the current learning rate
          cur_lr = list(optimizer.param_groups)[0]["lr"]
  
          step_time = time.time() - start_time
          epoch_time += step_time
  
          update_train_values = dict()
          loss_dict = dict()
          loss_dict["model_loss"] = loss.item()
          for key, value in loss_dict.items():
              update_train_values["avg_" + key] = value
          update_train_values["avg_loader_time"] = loader_time
          update_train_values["avg_step_time"] = step_time
          keep_avg.update_values(update_train_values)
  
          # print training stats
          if global_step % c.print_step == 0:
              log_dict = {
                  "step_time": [step_time, 2],
                  "loader_time": [loader_time, 4],
                  "current_lr": cur_lr,
              }
              c_logger.print_train_step(
                  batch_n_iter,
                  num_iter,
                  global_step,
                  log_dict,
                  loss_dict,
                  keep_avg.avg_values,
              )
  
          # plot step stats
          if global_step % 10 == 0:
              iter_stats = {"lr": cur_lr, "step_time": step_time}
              iter_stats.update(loss_dict)
              tb_logger.tb_train_iter_stats(global_step, iter_stats)
  
          # save checkpoint
          if global_step % c.save_step == 0:
              if c.checkpoint:
                  # save model
                  save_checkpoint(
                      model,
                      optimizer,
                      scheduler,
                      None,
                      None,
                      None,
                      global_step,
                      epoch,
                      OUT_PATH,
                      model_losses=loss_dict,
                      scaler=scaler.state_dict()
                      if c.mixed_precision
                      else None,
                  )
  
              # synthesize a full voice
              rand_idx = random.randrange(0, len(train_data))
              wav_path = (
                  train_data[rand_idx]
                  if not isinstance(train_data[rand_idx], (tuple, list))
                  else train_data[rand_idx][0]
              )
              wav = ap.load_wav(wav_path)
              ground_mel = ap.melspectrogram(wav)
              sample_wav = model.generate(
                  ground_mel,
                  c.batched,
                  c.target_samples,
                  c.overlap_samples,
                  use_cuda,
              )
              predict_mel = ap.melspectrogram(sample_wav)
  
              # compute spectrograms
              figures = {
                  "train/ground_truth": plot_spectrogram(ground_mel.T),
                  "train/prediction": plot_spectrogram(predict_mel.T),
              }
              tb_logger.tb_train_figures(global_step, figures)
  
              # Sample audio
              tb_logger.tb_train_audios(
                  global_step,
                  {"train/audio": sample_wav},
                  c.audio["sample_rate"],
              )
          end_time = time.time()
  
      # print epoch stats
      c_logger.print_train_epoch_end(
          global_step, epoch, epoch_time, keep_avg
      )
  
      # Plot Training Epoch Stats
      epoch_stats = {"epoch_time": epoch_time}
      epoch_stats.update(keep_avg.avg_values)
      tb_logger.tb_train_epoch_stats(global_step, epoch_stats)
      # TODO: plot model stats
      # if c.tb_model_param_stats:
      # tb_logger.tb_model_weights(model, global_step)
      return keep_avg.avg_values, global_step
  
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 3 for calling a global function with example usage: train(model, optimizer, criterion, scheduler, scaler, ap, global_step, epoch) and returns: keep_avg,  global_step