code: def main(args):
      # Load graph data.
      if args.dataset == "aifb":
          dataset = AIFBDataset()
      elif args.dataset == "bgs":
          dataset = BGSDataset()
      else:
          raise ValueError()
  
      g = dataset[0]
      category = dataset.predict_category
      num_classes = dataset.num_classes
      train_mask = g.nodes[category].data.pop("train_mask")
      test_mask = g.nodes[category].data.pop("test_mask")
      train_idx = th.nonzero(train_mask, as_tuple=False).squeeze()
      test_idx = th.nonzero(test_mask, as_tuple=False).squeeze()
      labels = g.nodes[category].data.pop("labels")
  
      # Split dataset into train, validate, test.
      val_idx = train_idx[: len(train_idx) // 5]
      train_idx = train_idx[len(train_idx) // 5 :]
  
      embed_layer = RelGraphEmbed(
          {ntype: g.num_nodes(ntype) for ntype in g.ntypes}, 16
      )
  
      # Create model.
      model = EntityClassify(
          16,
          num_classes,
          list(set(g.etypes)),
          embed_layer,
      )
  
      # Optimizer.
      optimizer = th.optim.Adam(
          model.parameters(), lr=1e-2, weight_decay=0
      )
  
      # Construct hetero sparse matrix.
      A = {}
      for stype, etype, dtype in g.canonical_etypes:
          eg = g[stype, etype, dtype]
          indices = th.stack(eg.edges("uv"))
          A[(stype, etype, dtype)] = dglsp.spmatrix(
              indices, shape=(g.num_nodes(stype), g.num_nodes(dtype))
          )
          ###########################################################
          # (HIGHLIGHT) Compute the normalized adjacency matrix with
          # Sparse Matrix API
          ###########################################################
          D1_hat = dglsp.diag(A[(stype, etype, dtype)].sum(1)) ** -0.5
          D2_hat = dglsp.diag(A[(stype, etype, dtype)].sum(0)) ** -0.5
          A[(stype, etype, dtype)] = (
              D1_hat @ A[(stype, etype, dtype)] @ D2_hat
          )
  
      # Training loop.
      print("start training...")
      model.train()
      for epoch in range(10):
          optimizer.zero_grad()
          logits = model(A)[category]
          loss = F.cross_entropy(logits[train_idx], labels[train_idx])
          loss.backward()
          optimizer.step()
  
          train_acc = th.sum(
              logits[train_idx].argmax(dim=1) == labels[train_idx]
          ).item() / len(train_idx)
          val_loss = F.cross_entropy(logits[val_idx], labels[val_idx])
          val_acc = th.sum(
              logits[val_idx].argmax(dim=1) == labels[val_idx]
          ).item() / len(val_idx)
          print(
              f"Epoch {epoch:05d} | Train Acc: {train_acc:.4f} | "
              f"Train Loss: {loss.item():.4f} | Valid Acc: {val_acc:.4f} | "
              f"Valid loss: {val_loss.item():.4f} "
          )
      print()
  
      model.eval()
      logits = model.forward(A)[category]
      test_loss = F.cross_entropy(logits[test_idx], labels[test_idx])
      test_acc = th.sum(
          logits[test_idx].argmax(dim=1) == labels[test_idx]
      ).item() / len(test_idx)
      print(
          "Test Acc: {:.4f} | Test loss: {:.4f}".format(
              test_acc, test_loss.item()
          )
      )
      print()
  
imports: import argparse
  import time
  import dgl
  import dgl.sparse
  import numpy
  import torch
  import torch.nn
  import torch.nn.functional
  from dgl.data.rdf import AIFBDataset, AMDataset, BGSDataset, MUTAGDataset
label_desc: Use image node_id 1 for calling a global function with example usage: main(args) without return types