code: def train(args):
      """Train E2E-TTS model."""
      set_deterministic_pytorch(args)
  
      # check cuda availability
      if not torch.cuda.is_available():
          logging.warning("cuda is not available")
  
      # get input and output dimension info
      with open(args.valid_json, "rb") as f:
          valid_json = json.load(f)["utts"]
      utts = list(valid_json.keys())
  
      # reverse input and output dimension
      idim = int(valid_json[utts[0]]["output"][0]["shape"][1])
      odim = int(valid_json[utts[0]]["input"][0]["shape"][1])
      logging.info("#input dims : " + str(idim))
      logging.info("#output dims: " + str(odim))
  
      # get extra input and output dimenstion
      if args.use_speaker_embedding:
          args.spk_embed_dim = int(
              valid_json[utts[0]]["input"][1]["shape"][0]
          )
      else:
          args.spk_embed_dim = None
      if args.use_second_target:
          args.spc_dim = int(
              valid_json[utts[0]]["input"][1]["shape"][1]
          )
      else:
          args.spc_dim = None
  
      # write model config
      if not os.path.exists(args.outdir):
          os.makedirs(args.outdir)
      model_conf = args.outdir + "/model.json"
      with open(model_conf, "wb") as f:
          logging.info("writing a model config file to" + model_conf)
          f.write(
              json.dumps(
                  (idim, odim, vars(args)),
                  indent=4,
                  ensure_ascii=False,
                  sort_keys=True,
              ).encode("utf_8")
          )
      for key in sorted(vars(args).keys()):
          logging.info("ARGS: " + key + ": " + str(vars(args)[key]))
  
      # specify model architecture
      if args.enc_init is not None or args.dec_init is not None:
          model = load_trained_modules(idim, odim, args, TTSInterface)
      else:
          model_class = dynamic_import(args.model_module)
          model = model_class(idim, odim, args)
      assert isinstance(model, TTSInterface)
      logging.info(model)
      reporter = model.reporter
  
      # check the use of multi-gpu
      if args.ngpu > 1:
          model = torch.nn.DataParallel(
              model, device_ids=list(range(args.ngpu))
          )
          if args.batch_size != 0:
              logging.warning(
                  "batch size is automatically increased (%d -> %d)"
                  % (args.batch_size, args.batch_size * args.ngpu)
              )
              args.batch_size *= args.ngpu
  
      # set torch device
      device = torch.device("cuda" if args.ngpu > 0 else "cpu")
      model = model.to(device)
  
      # freeze modules, if specified
      if args.freeze_mods:
          if hasattr(model, "module"):
              freeze_mods = ["module." + x for x in args.freeze_mods]
          else:
              freeze_mods = args.freeze_mods
  
          for mod, param in model.named_parameters():
              if any(mod.startswith(key) for key in freeze_mods):
                  logging.info(f"{mod} is frozen not to be updated.")
                  param.requires_grad = False
  
          model_params = filter(
              lambda x: x.requires_grad, model.parameters()
          )
      else:
          model_params = model.parameters()
  
      logging.warning(
          "num. model params: {:,} (num. trained: {:,} ({:.1f}%))".format(
              sum(p.numel() for p in model.parameters()),
              sum(
                  p.numel()
                  for p in model.parameters()
                  if p.requires_grad
              ),
              sum(
                  p.numel()
                  for p in model.parameters()
                  if p.requires_grad
              )
              * 100.0
              / sum(p.numel() for p in model.parameters()),
          )
      )
  
      # Setup an optimizer
      if args.opt == "adam":
          optimizer = torch.optim.Adam(
              model_params,
              args.lr,
              eps=args.eps,
              weight_decay=args.weight_decay,
          )
      elif args.opt == "noam":
          from espnet.nets.pytorch_backend.transformer.optimizer import (
              get_std_opt,
          )
  
          optimizer = get_std_opt(
              model_params,
              args.adim,
              args.transformer_warmup_steps,
              args.transformer_lr,
          )
      else:
          raise NotImplementedError("unknown optimizer: " + args.opt)
  
      # FIXME: TOO DIRTY HACK
      setattr(optimizer, "target", reporter)
      setattr(optimizer, "serialize", lambda s: reporter.serialize(s))
  
      # read json data
      with open(args.train_json, "rb") as f:
          train_json = json.load(f)["utts"]
      with open(args.valid_json, "rb") as f:
          valid_json = json.load(f)["utts"]
  
      use_sortagrad = args.sortagrad == -1 or args.sortagrad > 0
      if use_sortagrad:
          args.batch_sort_key = "input"
      # make minibatch list (variable length)
      train_batchset = make_batchset(
          train_json,
          args.batch_size,
          args.maxlen_in,
          args.maxlen_out,
          args.minibatches,
          batch_sort_key=args.batch_sort_key,
          min_batch_size=args.ngpu if args.ngpu > 1 else 1,
          shortest_first=use_sortagrad,
          count=args.batch_count,
          batch_bins=args.batch_bins,
          batch_frames_in=args.batch_frames_in,
          batch_frames_out=args.batch_frames_out,
          batch_frames_inout=args.batch_frames_inout,
          swap_io=True,
          iaxis=0,
          oaxis=0,
      )
      valid_batchset = make_batchset(
          valid_json,
          args.batch_size,
          args.maxlen_in,
          args.maxlen_out,
          args.minibatches,
          batch_sort_key=args.batch_sort_key,
          min_batch_size=args.ngpu if args.ngpu > 1 else 1,
          count=args.batch_count,
          batch_bins=args.batch_bins,
          batch_frames_in=args.batch_frames_in,
          batch_frames_out=args.batch_frames_out,
          batch_frames_inout=args.batch_frames_inout,
          swap_io=True,
          iaxis=0,
          oaxis=0,
      )
  
      load_tr = LoadInputsAndTargets(
          mode="tts",
          use_speaker_embedding=args.use_speaker_embedding,
          use_second_target=args.use_second_target,
          preprocess_conf=args.preprocess_conf,
          preprocess_args={
              "train": True
          },  # Switch the mode of preprocessing
          keep_all_data_on_mem=args.keep_all_data_on_mem,
      )
  
      load_cv = LoadInputsAndTargets(
          mode="tts",
          use_speaker_embedding=args.use_speaker_embedding,
          use_second_target=args.use_second_target,
          preprocess_conf=args.preprocess_conf,
          preprocess_args={
              "train": False
          },  # Switch the mode of preprocessing
          keep_all_data_on_mem=args.keep_all_data_on_mem,
      )
  
      converter = CustomConverter()
      # hack to make batchsize argument as 1
      # actual bathsize is included in a list
      train_iter = {
          "main": ChainerDataLoader(
              dataset=TransformDataset(
                  train_batchset,
                  lambda data: converter([load_tr(data)]),
              ),
              batch_size=1,
              num_workers=args.num_iter_processes,
              shuffle=not use_sortagrad,
              collate_fn=lambda x: x[0],
          )
      }
      valid_iter = {
          "main": ChainerDataLoader(
              dataset=TransformDataset(
                  valid_batchset,
                  lambda data: converter([load_cv(data)]),
              ),
              batch_size=1,
              shuffle=False,
              collate_fn=lambda x: x[0],
              num_workers=args.num_iter_processes,
          )
      }
  
      # Set up a trainer
      updater = CustomUpdater(
          model,
          args.grad_clip,
          train_iter,
          optimizer,
          device,
          args.accum_grad,
      )
      trainer = training.Trainer(
          updater, (args.epochs, "epoch"), out=args.outdir
      )
  
      # Resume from a snapshot
      if args.resume:
          logging.info("resumed from %s" % args.resume)
          torch_resume(args.resume, trainer)
  
      # set intervals
      eval_interval = (args.eval_interval_epochs, "epoch")
      save_interval = (args.save_interval_epochs, "epoch")
      report_interval = (args.report_interval_iters, "iteration")
  
      # Evaluate the model with the test dataset for each epoch
      trainer.extend(
          CustomEvaluator(model, valid_iter, reporter, device),
          trigger=eval_interval,
      )
  
      # Save snapshot for each epoch
      trainer.extend(torch_snapshot(), trigger=save_interval)
  
      # Save best models
      trainer.extend(
          snapshot_object(model, "model.loss.best"),
          trigger=training.triggers.MinValueTrigger(
              "validation/main/loss", trigger=eval_interval
          ),
      )
  
      # Save attention figure for each epoch
      if args.num_save_attention > 0:
          data = sorted(
              list(valid_json.items())[: args.num_save_attention],
              key=lambda x: int(x[1]["output"][0]["shape"][0]),
          )
          if hasattr(model, "module"):
              att_vis_fn = model.module.calculate_all_attentions
              plot_class = model.module.attention_plot_class
              reduction_factor = model.module.reduction_factor
          else:
              att_vis_fn = model.calculate_all_attentions
              plot_class = model.attention_plot_class
              reduction_factor = model.reduction_factor
          if reduction_factor > 1:
              # fix the length to crop attention weight plot correctly
              data = copy.deepcopy(data)
              for idx in range(len(data)):
                  ilen = data[idx][1]["input"][0]["shape"][0]
                  data[idx][1]["input"][0]["shape"][0] = (
                      ilen // reduction_factor
                  )
          att_reporter = plot_class(
              att_vis_fn,
              data,
              args.outdir + "/att_ws",
              converter=converter,
              transform=load_cv,
              device=device,
              reverse=True,
          )
          trainer.extend(att_reporter, trigger=eval_interval)
      else:
          att_reporter = None
  
      # Make a plot for training and validation values
      if hasattr(model, "module"):
          base_plot_keys = model.module.base_plot_keys
      else:
          base_plot_keys = model.base_plot_keys
      plot_keys = []
      for key in base_plot_keys:
          plot_key = ["main/" + key, "validation/main/" + key]
          trainer.extend(
              extensions.PlotReport(
                  plot_key, "epoch", file_name=key + ".png"
              ),
              trigger=eval_interval,
          )
          plot_keys += plot_key
      trainer.extend(
          extensions.PlotReport(
              plot_keys, "epoch", file_name="all_loss.png"
          ),
          trigger=eval_interval,
      )
  
      # Write a log of evaluation statistics for each epoch
      trainer.extend(extensions.LogReport(trigger=report_interval))
      report_keys = ["epoch", "iteration", "elapsed_time"] + plot_keys
      trainer.extend(
          extensions.PrintReport(report_keys), trigger=report_interval
      )
      trainer.extend(extensions.ProgressBar(), trigger=report_interval)
  
      set_early_stop(trainer, args)
      if (
          args.tensorboard_dir is not None
          and args.tensorboard_dir != ""
      ):
          from torch.utils.tensorboard import SummaryWriter
  
          writer = SummaryWriter(args.tensorboard_dir)
          trainer.extend(
              TensorboardLogger(writer, att_reporter),
              trigger=report_interval,
          )
  
      if use_sortagrad:
          trainer.extend(
              ShufflingEnabler([train_iter]),
              trigger=(
                  args.sortagrad
                  if args.sortagrad != -1
                  else args.epochs,
                  "epoch",
              ),
          )
  
      # Run the training
      trainer.run()
      check_early_stop(trainer, args.epochs)
  
imports: import copy
  import json
  import logging
  import math
  import os
  import time
  import chainer
  import kaldiio
  import numpy
  import torch
  from chainer import training
  from chainer.training import extensions
  from espnet.asr.asr_utils import get_model_conf, snapshot_object, torch_load, torch_resume, torch_snapshot
  from espnet.asr.pytorch_backend.asr_init import load_trained_modules
  from espnet.nets.pytorch_backend.nets_utils import pad_list
  from espnet.nets.tts_interface import TTSInterface
  from espnet.utils.dataset import ChainerDataLoader, TransformDataset
  from espnet.utils.deterministic_utils import set_deterministic_pytorch
  from espnet.utils.dynamic_import import dynamic_import
  from espnet.utils.io_utils import LoadInputsAndTargets
  from espnet.utils.training.batchfy import make_batchset
  from espnet.utils.training.evaluator import BaseEvaluator
  from espnet.utils.training.iterators import ShufflingEnabler
  from espnet.utils.training.tensorboard_logger import TensorboardLogger
  from espnet.utils.training.train_utils import check_early_stop, set_early_stop
label_desc: Use image node_id 1 for calling a global function with example usage: train(args) without return types