code: def run_epoch(
      epoch,
      data_iter,
      dev_rank,
      ndev,
      model,
      loss_compute,
      is_train=True,
  ):
      universal = isinstance(model, UTransformer)
      with loss_compute:
          for i, g in enumerate(data_iter):
              with T.set_grad_enabled(is_train):
                  if universal:
                      output, loss_act = model(g)
                      if is_train:
                          loss_act.backward(retain_graph=True)
                  else:
                      output = model(g)
                  tgt_y = g.tgt_y
                  n_tokens = g.n_tokens
                  loss = loss_compute(output, tgt_y, n_tokens)
  
      if universal:
          for step in range(1, model.MAX_DEPTH + 1):
              print(
                  "nodes entering step {}: {:.2f}%".format(
                      step, (1.0 * model.stat[step] / model.stat[0])
                  )
              )
          model.reset_stat()
      print(
          "Epoch {} {}: Dev {} average loss: {}, accuracy {}".format(
              epoch,
              "Training" if is_train else "Evaluating",
              dev_rank,
              loss_compute.avg_loss,
              loss_compute.accuracy,
          )
      )
  
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 1 for calling a global function with example usage: run_epoch(epoch, data_iter, dev_rank, ndev, model, loss_compute, is_train) without return types