code: def main(dev_id, args):
      if dev_id == -1:
          device = torch.device("cpu")
      else:
          device = torch.device("cuda:{}".format(dev_id))
      # Set current device
      th.cuda.set_device(device)
      # Prepare dataset
      dataset = get_dataset(args.dataset)
      V = dataset.vocab_size
      criterion = LabelSmoothing(
          V, padding_idx=dataset.pad_id, smoothing=0.1
      )
      dim_model = 512
      # Build graph pool
      graph_pool = GraphPool()
      # Create model
      model = make_model(
          V, V, N=args.N, dim_model=dim_model, universal=args.universal
      )
      # Sharing weights between Encoder & Decoder
      model.src_embed.lut.weight = model.tgt_embed.lut.weight
      model.generator.proj.weight = model.tgt_embed.lut.weight
      # Move model to corresponding device
      model, criterion = model.to(device), criterion.to(device)
      # Loss function
      if args.ngpu > 1:
          dev_rank = dev_id  # current device id
          ndev = args.ngpu  # number of devices (including cpu)
          loss_compute = partial(
              MultiGPULossCompute,
              criterion,
              args.ngpu,
              args.grad_accum,
              model,
          )
      else:  # cpu or single gpu case
          dev_rank = 0
          ndev = 1
          loss_compute = partial(
              SimpleLossCompute, criterion, args.grad_accum
          )
  
      if ndev > 1:
          for param in model.parameters():
              dist.all_reduce(param.data, op=dist.ReduceOp.SUM)
              param.data /= ndev
  
      # Optimizer
      model_opt = NoamOpt(
          dim_model,
          0.1,
          4000,
          T.optim.Adam(
              model.parameters(), lr=1e-3, betas=(0.9, 0.98), eps=1e-9
          ),
      )
  
      # Train & evaluate
      for epoch in range(100):
          start = time.time()
          train_iter = dataset(
              graph_pool,
              mode="train",
              batch_size=args.batch,
              device=device,
              dev_rank=dev_rank,
              ndev=ndev,
          )
          model.train(True)
          run_epoch(
              epoch,
              train_iter,
              dev_rank,
              ndev,
              model,
              loss_compute(opt=model_opt),
              is_train=True,
          )
          if dev_rank == 0:
              model.att_weight_map = None
              model.eval()
              valid_iter = dataset(
                  graph_pool,
                  mode="valid",
                  batch_size=args.batch,
                  device=device,
                  dev_rank=dev_rank,
                  ndev=1,
              )
              run_epoch(
                  epoch,
                  valid_iter,
                  dev_rank,
                  1,
                  model,
                  loss_compute(opt=None),
                  is_train=False,
              )
              end = time.time()
              print("epoch time: {}".format(end - start))
  
              # Visualize attention
              if args.viz:
                  src_seq = dataset.get_seq_by_id(
                      VIZ_IDX, mode="valid", field="src"
                  )
                  tgt_seq = dataset.get_seq_by_id(
                      VIZ_IDX, mode="valid", field="tgt"
                  )[:-1]
                  draw_atts(
                      model.att_weight_map,
                      src_seq,
                      tgt_seq,
                      exp_setting,
                      "epoch_{}".format(epoch),
                  )
              args_filter = [
                  "batch",
                  "gpus",
                  "viz",
                  "master_ip",
                  "master_port",
                  "grad_accum",
                  "ngpu",
              ]
              exp_setting = "-".join(
                  "{}".format(v)
                  for k, v in vars(args).items()
                  if k not in args_filter
              )
              with open(
                  "checkpoints/{}-{}.pkl".format(exp_setting, epoch),
                  "wb",
              ) as f:
                  torch.save(model.state_dict(), f)
  
file_path: dgl/examples/pytorch/transformer/translation_train.py
imports: import argparse
  from functools import partial
  import numpy
  import torch
  import torch.distributed
  from dataset import *
  from loss import *
  from modules import *
  from modules.config import *
  from optims import *
label_desc: Use image node_id 3 for calling a global function with example usage: main(dev_id, args) without return types