code: def main(args):
      # load and preprocess dataset
      if args.dataset == "cora":
          data = CoraGraphDataset()
      elif args.dataset == "citeseer":
          data = CiteseerGraphDataset()
      elif args.dataset == "pubmed":
          data = PubmedGraphDataset()
      else:
          raise ValueError("Unknown dataset: {}".format(args.dataset))
  
      g = data[0]
      if args.gpu < 0:
          cuda = False
          ctx = mx.cpu(0)
      else:
          cuda = True
          ctx = mx.gpu(args.gpu)
          g = g.to(ctx)
  
      features = g.ndata["feat"]
      labels = mx.nd.array(g.ndata["label"], dtype="float32", ctx=ctx)
      train_mask = g.ndata["train_mask"]
      val_mask = g.ndata["val_mask"]
      test_mask = g.ndata["test_mask"]
      in_feats = features.shape[1]
      n_classes = data.num_classes
      n_edges = data.graph.number_of_edges()
      print(
          """----Data statistics------'
        #Edges %d
        #Classes %d
        #Train samples %d
        #Val samples %d
        #Test samples %d"""
          % (
              n_edges,
              n_classes,
              train_mask.sum().asscalar(),
              val_mask.sum().asscalar(),
              test_mask.sum().asscalar(),
          )
      )
  
      # add self loop
      g = dgl.remove_self_loop(g)
      g = dgl.add_self_loop(g)
  
      # create TAGCN model
      model = TAGCN(
          g,
          in_feats,
          args.n_hidden,
          n_classes,
          args.n_layers,
          mx.nd.relu,
          args.dropout,
      )
  
      model.initialize(ctx=ctx)
      n_train_samples = train_mask.sum().asscalar()
      loss_fcn = gluon.loss.SoftmaxCELoss()
  
      # use optimizer
      print(model.collect_params())
      trainer = gluon.Trainer(
          model.collect_params(),
          "adam",
          {"learning_rate": args.lr, "wd": args.weight_decay},
      )
  
      # initialize graph
      dur = []
      for epoch in range(args.n_epochs):
          if epoch >= 3:
              t0 = time.time()
          # forward
          with mx.autograd.record():
              pred = model(features)
              loss = loss_fcn(
                  pred, labels, mx.nd.expand_dims(train_mask, 1)
              )
              loss = loss.sum() / n_train_samples
  
          loss.backward()
          trainer.step(batch_size=1)
  
          if epoch >= 3:
              loss.asscalar()
              dur.append(time.time() - t0)
              acc = evaluate(model, features, labels, val_mask)
              print(
                  "Epoch {:05d} | Time(s) {:.4f} | Loss {:.4f} | Accuracy {:.4f} | "
                  "ETputs(KTEPS) {:.2f}".format(
                      epoch,
                      np.mean(dur),
                      loss.asscalar(),
                      acc,
                      n_edges / np.mean(dur) / 1000,
                  )
              )
  
      print()
      acc = evaluate(model, features, labels, val_mask)
      print("Test accuracy {:.2%}".format(acc))
  
file_path: dgl/examples/mxnet/tagcn/train.py
imports: import argparse
  import time
  import dgl
  import mxnet
  import networkx
  import numpy
  from dgl.data import CiteseerGraphDataset, CoraGraphDataset, PubmedGraphDataset, register_data_args
  from mxnet import gluon
  from tagcn import TAGCN
label_desc: Use image node_id 2 for calling a global function with example usage: main(args) without return types