code: def main():
      parser = argparse.ArgumentParser(
          description="OGBL(Full Batch GCN/GraphSage + NGNN)"
      )
  
      # dataset setting
      parser.add_argument(
          "--dataset",
          type=str,
          default="ogbl-ddi",
          choices=["ogbl-ddi", "ogbl-collab", "ogbl-ppa"],
      )
  
      # device setting
      parser.add_argument(
          "--device",
          type=int,
          default=0,
          help="GPU device ID. Use -1 for CPU training.",
      )
  
      # model structure settings
      parser.add_argument(
          "--use_sage",
          action="store_true",
          help="If not set, use GCN by default.",
      )
      parser.add_argument(
          "--ngnn_type",
          type=str,
          default="input",
          choices=["input", "hidden"],
          help="You can set this value from 'input' or 'hidden' to apply NGNN to different GNN layers.",
      )
      parser.add_argument(
          "--num_layers",
          type=int,
          default=3,
          help="number of GNN layers",
      )
      parser.add_argument("--hidden_channels", type=int, default=256)
      parser.add_argument("--dropout", type=float, default=0.0)
      parser.add_argument("--batch_size", type=int, default=64 * 1024)
      parser.add_argument("--lr", type=float, default=0.001)
      parser.add_argument("--epochs", type=int, default=400)
  
      # training settings
      parser.add_argument("--eval_steps", type=int, default=1)
      parser.add_argument("--runs", type=int, default=10)
      args = parser.parse_args()
      print(args)
  
      device = (
          f"cuda:{args.device}"
          if args.device != -1 and torch.cuda.is_available()
          else "cpu"
      )
      device = torch.device(device)
  
      dataset = DglLinkPropPredDataset(name=args.dataset)
      g = dataset[0]
      split_edge = dataset.get_edge_split()
  
      # We randomly pick some training samples that we want to evaluate on:
      idx = torch.randperm(split_edge["train"]["edge"].size(0))
      idx = idx[: split_edge["valid"]["edge"].size(0)]
      split_edge["eval_train"] = {
          "edge": split_edge["train"]["edge"][idx]
      }
  
      if dataset.name == "ogbl-ppa":
          g.ndata["feat"] = g.ndata["feat"].to(torch.float)
  
      if dataset.name == "ogbl-ddi":
          emb = torch.nn.Embedding(
              g.num_nodes(), args.hidden_channels
          ).to(device)
          in_channels = args.hidden_channels
      else:  # ogbl-collab, ogbl-ppa
          in_channels = g.ndata["feat"].size(-1)
  
      # select model
      if args.use_sage:
          model = SAGE(
              in_channels,
              args.hidden_channels,
              args.hidden_channels,
              args.num_layers,
              args.dropout,
              args.ngnn_type,
              dataset.name,
          )
      else:  # GCN
          g = dgl.add_self_loop(g)
          model = GCN(
              in_channels,
              args.hidden_channels,
              args.hidden_channels,
              args.num_layers,
              args.dropout,
              args.ngnn_type,
              dataset.name,
          )
  
      predictor = LinkPredictor(
          args.hidden_channels, args.hidden_channels, 1, 3, args.dropout
      )
  
      g, model, predictor = map(
          lambda x: x.to(device), (g, model, predictor)
      )
  
      evaluator = Evaluator(name=dataset.name)
      loggers = {
          "Hits@20": Logger(args.runs, args),
          "Hits@50": Logger(args.runs, args),
          "Hits@100": Logger(args.runs, args),
      }
  
      for run in range(args.runs):
          model.reset_parameters()
          predictor.reset_parameters()
          if dataset.name == "ogbl-ddi":
              torch.nn.init.xavier_uniform_(emb.weight)
              g.ndata["feat"] = emb.weight
          optimizer = torch.optim.Adam(
              list(model.parameters())
              + list(predictor.parameters())
              + (
                  list(emb.parameters())
                  if dataset.name == "ogbl-ddi"
                  else []
              ),
              lr=args.lr,
          )
          for epoch in range(1, 1 + args.epochs):
              loss = train(
                  model,
                  predictor,
                  g,
                  g.ndata["feat"],
                  split_edge,
                  optimizer,
                  args.batch_size,
              )
  
              if epoch % args.eval_steps == 0:
                  results = test(
                      model,
                      predictor,
                      g,
                      g.ndata["feat"],
                      split_edge,
                      evaluator,
                      args.batch_size,
                  )
                  for key, result in results.items():
                      loggers[key].add_result(run, result)
                      train_hits, valid_hits, test_hits = result
                      print(key)
                      print(
                          f"Run: {run + 1:02d}, "
                          f"Epoch: {epoch:02d}, "
                          f"Loss: {loss:.4f}, "
                          f"Train: {100 * train_hits:.2f}%, "
                          f"Valid: {100 * valid_hits:.2f}%, "
                          f"Test: {100 * test_hits:.2f}%"
                      )
                  print("---")
  
          for key in loggers.keys():
              print(key)
              loggers[key].print_statistics(run)
  
      for key in loggers.keys():
          print(key)
          loggers[key].print_statistics()
  
imports: import argparse
  import math
  import dgl
  import torch
  import torch.nn.functional
  from dgl.dataloading.negative_sampler import GlobalUniform
  from dgl.nn.pytorch import GraphConv, SAGEConv
  from ogb.linkproppred import DglLinkPropPredDataset, Evaluator
  from torch.nn import Linear
  from torch.utils.data import DataLoader
label_desc: Use image node_id 3 for calling a global function with example usage: main() without return types