code: def train(model, criterion, optimizer, scheduler, ap, global_step):
      data_loader = setup_loader(ap, is_val=False, verbose=True)
      model.train()
      epoch_time = 0
      best_loss = float("inf")
      avg_loss = 0
      avg_loader_time = 0
      end_time = time.time()
      for _, data in enumerate(data_loader):
          start_time = time.time()
  
          # setup input data
          inputs = data[0]
          loader_time = time.time() - end_time
          global_step += 1
  
          # setup lr
          if c.lr_decay:
              scheduler.step()
          optimizer.zero_grad()
  
          # dispatch data to GPU
          if use_cuda:
              inputs = inputs.cuda(non_blocking=True)
              # labels = labels.cuda(non_blocking=True)
  
          # forward pass model
          outputs = model(inputs)
  
          # loss computation
          loss = criterion(
              outputs.view(
                  c.num_speakers_in_batch,
                  outputs.shape[0] // c.num_speakers_in_batch,
                  -1,
              )
          )
          loss.backward()
          grad_norm, _ = check_update(model, c.grad_clip)
          optimizer.step()
  
          step_time = time.time() - start_time
          epoch_time += step_time
  
          # Averaged Loss and Averaged Loader Time
          avg_loss = (
              0.01 * loss.item() + 0.99 * avg_loss
              if avg_loss != 0
              else loss.item()
          )
          avg_loader_time = (
              1 / c.num_loader_workers * loader_time
              + (c.num_loader_workers - 1)
              / c.num_loader_workers
              * avg_loader_time
              if avg_loader_time != 0
              else loader_time
          )
          current_lr = optimizer.param_groups[0]["lr"]
  
          if global_step % c.steps_plot_stats == 0:
              # Plot Training Epoch Stats
              train_stats = {
                  "loss": avg_loss,
                  "lr": current_lr,
                  "grad_norm": grad_norm,
                  "step_time": step_time,
                  "avg_loader_time": avg_loader_time,
              }
              tb_logger.tb_train_epoch_stats(global_step, train_stats)
              figures = {
                  # FIXME: not constant
                  "UMAP Plot": plot_embeddings(
                      outputs.detach().cpu().numpy(), 10
                  ),
              }
              tb_logger.tb_train_figures(global_step, figures)
  
          if global_step % c.print_step == 0:
              print(
                  "   | > Step:{}  Loss:{:.5f}  AvgLoss:{:.5f}  GradNorm:{:.5f}  "
                  "StepTime:{:.2f}  LoaderTime:{:.2f}  AvGLoaderTime:{:.2f}  LR:{:.6f}".format(
                      global_step,
                      loss.item(),
                      avg_loss,
                      grad_norm,
                      step_time,
                      loader_time,
                      avg_loader_time,
                      current_lr,
                  ),
                  flush=True,
              )
  
          # save best model
          best_loss = save_best_model(
              model,
              optimizer,
              avg_loss,
              best_loss,
              OUT_PATH,
              global_step,
          )
  
          end_time = time.time()
      return avg_loss, global_step
  
imports: import argparse
  import os
  import sys
  import time
  import traceback
  import torch
  from torch.utils.data import DataLoader
  from TTS.speaker_encoder.dataset import MyDataset
  from TTS.speaker_encoder.losses import AngleProtoLoss, GE2ELoss
  from TTS.speaker_encoder.model import SpeakerEncoder
  from TTS.speaker_encoder.utils.generic_utils import check_config_speaker_encoder, save_best_model
  from TTS.speaker_encoder.utils.visual import plot_embeddings
  from TTS.tts.datasets.preprocess import load_meta_data
  from TTS.utils.audio import AudioProcessor
  from TTS.utils.generic_utils import 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 NoamLR, check_update
label_desc: Use image node_id 2 for calling a global function with example usage: train(model, criterion, optimizer, scheduler, ap, global_step) and returns: avg_loss,  global_step