code: def train(args):
      (
          store,
          num_workers,
          rank,
          local_rank,
          is_master_node,
          ctx_l,
      ) = init_comm(args.comm_backend, args.gpus)
      task = get_task(args.task_name, args.train_dir, args.eval_dir)
      # setup_logging(args, local_rank)
      # random seed
      set_seed(args.seed)
      level = logging.INFO
      detail_dir = os.path.join(args.output_dir, args.task_name)
      if not os.path.exists(detail_dir):
          os.mkdir(detail_dir)
      logging_config(
          detail_dir,
          name="train_{}_{}_".format(args.task_name, args.model_name)
          + str(rank),  # avoid race
          level=level,
          console=(local_rank == 0),
      )
      logging.info(args)
      cfg, tokenizer, classify_net, use_segmentation = get_network(
          args.model_name,
          ctx_l,
          args.param_checkpoint,
          args.backbone_path,
          task,
      )
  
      logging.info("Prepare training data")
      train_data, _ = get_task_data(
          args, task, tokenizer, segment="train"
      )
      train_batchify = bf.Group(
          bf.Group(bf.Pad(), bf.Pad(), bf.Stack()), bf.Stack()
      )
  
      rs = np.random.RandomState(100)
      rs.shuffle(train_data)
      sampler = SplitSampler(
          len(train_data),
          num_parts=num_workers,
          part_index=rank,
          even_size=True,
      )
  
      dataloader = DataLoader(
          train_data,
          batch_size=args.batch_size,
          batchify_fn=train_batchify,
          num_workers=0,
          sampler=sampler,
      )
  
      param_dict = classify_net.collect_params()
      # Do not apply weight decay to all the LayerNorm and bias
      for _, v in classify_net.collect_params(
          ".*beta|.*gamma|.*bias"
      ).items():
          v.wd_mult = 0.0
      # Set grad_req if gradient accumulation is required
      params = [p for p in param_dict.values() if p.grad_req != "null"]
      num_accumulated = args.num_accumulated
      if num_accumulated > 1:
          logging.info(
              "Using gradient accumulation. Effective global batch size = {}".format(
                  num_accumulated
                  * args.batch_size
                  * len(ctx_l)
                  * num_workers
              )
          )
          for p in params:
              p.grad_req = "add"
      if local_rank == 0:
          writer = SummaryWriter(
              logdir=os.path.join(
                  args.output_dir,
                  args.task_name
                  + "_tensorboard_"
                  + str(args.lr)
                  + "_"
                  + str(args.epochs),
              )
          )
      if args.comm_backend == "horovod":
          # Horovod: fetch and broadcast parameters
          hvd.broadcast_parameters(param_dict, root_rank=0)
  
      epoch_size = (len(dataloader) + len(ctx_l) - 1) // len(ctx_l)
      max_update = epoch_size * args.epochs
      warmup_steps = int(np.ceil(max_update * args.warmup_ratio))
  
      dataloader = grouper(repeat(dataloader), len(ctx_l))
  
      lr_scheduler = PolyScheduler(
          max_update=max_update,
          base_lr=args.lr,
          warmup_begin_lr=0.0,
          pwr=1,
          final_lr=0.0,
          warmup_steps=warmup_steps,
          warmup_mode="linear",
      )
      optimizer_params = {
          "learning_rate": args.lr,
          "wd": args.wd,
          "lr_scheduler": lr_scheduler,
      }
      if args.comm_backend == "horovod":
          trainer = hvd.DistributedTrainer(
              param_dict, args.optimizer, optimizer_params
          )
      else:
          trainer = mx.gluon.Trainer(
              classify_net.collect_params(), "adamw", optimizer_params
          )
  
      if args.task_name == "sts":
          loss_function = gluon.loss.L2Loss()
      else:
          loss_function = gluon.loss.SoftmaxCELoss()
  
      metrics = task.metric
      # prepare loss function
      log_loss = 0
      log_gnorm = 0
      log_step = 0
      if args.log_interval > 0:
          log_interval = args.log_interval
      else:
          log_interval = int(epoch_size * 0.5)
  
      start_time = time.time()
      total_loss = 0
      total_grad = 0
      total_step = 0
      for i in range(max_update):
          sample_l = next(dataloader)
          loss_l = []
          for sample, ctx in zip(sample_l, ctx_l):
              (token_ids, token_types, valid_length), label = sample
              # Move to the corresponding context
              token_ids = mx.np.array(token_ids, ctx=ctx)
              token_types = mx.np.array(token_types, ctx=ctx)
              valid_length = mx.np.array(valid_length, ctx=ctx)
              label = mx.np.array(label, ctx=ctx)
              with mx.autograd.record():
                  scores = classify_net(
                      token_ids, token_types, valid_length
                  )
                  loss = loss_function(scores, label).mean() / len(
                      ctx_l
                  )
                  loss_l.append(loss)
              if task.task_name == "sts":
                  label = label.reshape((-1, 1))
              for metric in metrics:
                  metric.update([label], [scores])
  
          for loss in loss_l:
              loss.backward()
          trainer.allreduce_grads()
          # Begin Norm Clipping
          total_norm, ratio, is_finite = clip_grad_global_norm(
              params, args.max_grad_norm
          )
          trainer.update(1.0)
          step_loss = sum([loss.asnumpy() for loss in loss_l])
          log_loss += step_loss
          log_gnorm += total_norm
          log_step += 1
          total_step += 1
          total_loss += step_loss
          total_grad += total_norm
          if local_rank == 0:
              writer.add_scalar(
                  "train_loss_avg", total_loss * 1.0 / total_step, i
              )
              writer.add_scalar("lr", trainer.learning_rate, i)
              writer.add_scalar("train_loss", step_loss, i)
              writer.add_scalar(
                  "grad_norm_avg", total_grad * 1.0 / total_step, i
              )
              writer.add_scalar("grad_norm", total_norm, i)
              for metric in metrics:
                  metric_name, result = metric.get()
                  writer.add_scalar(metric_name, result, i)
          if log_step >= log_interval or i == max_update - 1:
              curr_time = time.time()
              metric_log = ""
              for metric in metrics:
                  metric_nm, val = metric.get()
                  metric_log += ", {}: = {}".format(metric_nm, val)
              logging.info(
                  "[Iter {} / {}] avg {} = {:.2f}, avg gradient norm = {:.2f}, lr = {}, ETA={:.2f}h".format(
                      i + 1,
                      max_update,
                      "loss",
                      log_loss / log_step,
                      log_gnorm / log_step,
                      trainer.learning_rate,
                      (max_update - i)
                      * ((curr_time - start_time) / i)
                      / 3600,
                  )
                  + metric_log
              )
              log_loss = 0
              log_gnorm = 0
              log_step = 0
          if local_rank == 0 and (
              i == max_update - 1
              or i % (max_update // args.epochs) == 0
              and i > 0
          ):
              ckpt_name = "{}_{}_{}.params".format(
                  args.model_name, args.task_name, (i + 1)
              )
  
              params_saved = os.path.join(detail_dir, ckpt_name)
              classify_net.save_parameters(params_saved)
              logging.info("Params saved in: {}".format(params_saved))
              for metric in metrics:
                  metric.reset()
  
imports: import gluonnlp
  from tensorboardX import SummaryWriter
  import numpy
  import mxnet
  import json
  import random
  import pandas
  import os
  import logging
  import time
  import argparse
  import copy
  from mxnet.gluon.metric import Accuracy, F1, MCC, PearsonCorrelation, CompositeEvalMetric
  from classification_utils import get_task
  import matplotlib.pyplot
  from tqdm import tqdm
  from mxnet import gluon
  from gluonnlp.data.sampler import SplitSampler
  from mxnet.gluon import nn
  from gluonnlp.models import get_backbone
  from gluonnlp.utils.parameter import clip_grad_global_norm, count_parameters, deduplicate_param_dict
  from gluonnlp.utils.preprocessing import get_trimmed_lengths
  from gluonnlp.utils.misc import get_mxnet_visible_ctx, grouper, repeat, logging_config
  from mxnet.gluon.data import batchify
  from mxnet.gluon.data import DataLoader
  from mxnet.lr_scheduler import PolyScheduler
  from gluonnlp.utils import set_seed
  from gluonnlp.utils.misc import init_comm, parse_ctx
  from classification import TextPredictionNet
label_desc: Use image node_id 6 for calling a global function with example usage: train(args) without return types