code: def train(args):
      args = _parse_args(args)
      _LG.info("%s", args)
  
      args.save_dir.mkdir(parents=True, exist_ok=True)
      if "sox_io" in torchaudio.list_audio_backends():
          torchaudio.set_audio_backend("sox_io")
  
      start_epoch = 1
      if args.resume:
          checkpoint = torch.load(args.resume)
          if args.sample_rate != checkpoint["sample_rate"]:
              raise ValueError(
                  "The provided sample rate ({args.sample_rate}) does not match "
                  "the sample rate from the check point ({checkpoint['sample_rate']})."
              )
          if args.num_speakers != checkpoint["num_speakers"]:
              raise ValueError(
                  "The provided #of speakers ({args.num_speakers}) does not match "
                  "the #of speakers from the check point ({checkpoint['num_speakers']}.)"
              )
          start_epoch = checkpoint["epoch"]
  
      device = torch.device(
          "cuda" if torch.cuda.is_available() else "cpu"
      )
      _LG.info("Using: %s", device)
  
      model = _get_model(num_sources=args.num_speakers)
      model.to(device)
  
      model = torch.nn.parallel.DistributedDataParallel(
          model,
          device_ids=[device] if torch.cuda.is_available() else None,
      )
      optimizer = torch.optim.Adam(
          model.parameters(), lr=args.learning_rate
      )
  
      if args.resume:
          _LG.info("Loading parameters from the checkpoint...")
          model.module.load_state_dict(checkpoint["model"])
          optimizer.load_state_dict(checkpoint["optimizer"])
      else:
          dist_utils.synchronize_params(
              str(args.save_dir / "tmp.pt"), device, model, optimizer
          )
  
      lr_scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(
          optimizer, mode="max", factor=0.5, patience=3
      )
  
      train_loader, valid_loader, eval_loader = _get_dataloader(
          args.dataset,
          args.dataset_dir,
          args.num_speakers,
          args.sample_rate,
          args.batch_size,
      )
  
      num_train_samples = len(train_loader.dataset)
      num_valid_samples = len(valid_loader.dataset)
      num_eval_samples = len(eval_loader.dataset)
  
      _LG.info_on_master("Datasets:")
      _LG.info_on_master(" - Train: %s", num_train_samples)
      _LG.info_on_master(" - Valid: %s", num_valid_samples)
      _LG.info_on_master(" -  Eval: %s", num_eval_samples)
  
      trainer = conv_tasnet.trainer.Trainer(
          model,
          optimizer,
          train_loader,
          valid_loader,
          eval_loader,
          args.grad_clip,
          device,
          debug=args.debug,
      )
  
      log_path = args.save_dir / "log.csv"
      _write_header(log_path, args)
      dist_utils.write_csv_on_master(
          log_path,
          [
              "epoch",
              "learning_rate",
              "valid_si_snri",
              "valid_sdri",
              "eval_si_snri",
              "eval_sdri",
          ],
      )
  
      _LG.info_on_master("Running %s epochs", args.epochs)
      for epoch in range(start_epoch, start_epoch + args.epochs):
          _LG.info_on_master("=" * 70)
          _LG.info_on_master("Epoch: %s", epoch)
          _LG.info_on_master(
              "Learning rate: %s", optimizer.param_groups[0]["lr"]
          )
          _LG.info_on_master("=" * 70)
  
          t0 = time.monotonic()
          trainer.train_one_epoch()
          train_sps = num_train_samples / (time.monotonic() - t0)
  
          _LG.info_on_master("-" * 70)
  
          t0 = time.monotonic()
          valid_metric = trainer.validate()
          valid_sps = num_valid_samples / (time.monotonic() - t0)
          _LG.info_on_master("Valid: %s", valid_metric)
  
          _LG.info_on_master("-" * 70)
  
          t0 = time.monotonic()
          eval_metric = trainer.evaluate()
          eval_sps = num_eval_samples / (time.monotonic() - t0)
          _LG.info_on_master(" Eval: %s", eval_metric)
  
          _LG.info_on_master("-" * 70)
  
          _LG.info_on_master(
              "Train: Speed: %6.2f [samples/sec]", train_sps
          )
          _LG.info_on_master(
              "Valid: Speed: %6.2f [samples/sec]", valid_sps
          )
          _LG.info_on_master(
              " Eval: Speed: %6.2f [samples/sec]", eval_sps
          )
  
          _LG.info_on_master("-" * 70)
  
          dist_utils.write_csv_on_master(
              log_path,
              [
                  epoch,
                  optimizer.param_groups[0]["lr"],
                  valid_metric.si_snri,
                  valid_metric.sdri,
                  eval_metric.si_snri,
                  eval_metric.sdri,
              ],
          )
  
          lr_scheduler.step(valid_metric.si_snri)
  
          save_path = args.save_dir / f"epoch_{epoch}.pt"
          dist_utils.save_on_master(
              save_path,
              {
                  "model": model.module.state_dict(),
                  "optimizer": optimizer.state_dict(),
                  "num_speakers": args.num_speakers,
                  "sample_rate": args.sample_rate,
                  "epoch": epoch,
              },
          )
  
file_path: audio/examples/source_separation/conv_tasnet/train.py
imports: import argparse
  import pathlib
  import time
  import conv_tasnet
  import torch
  import torchaudio
  import torchaudio.models
  from utils import dist_utils
  from utils.dataset import utils
label_desc: Use image node_id 5 for calling a global function with example usage: train(args) without return types