code: def main():
      parser = argparse.ArgumentParser(
          description="OGBL-Citation2 (GraphSAINT)"
      )
      parser.add_argument("--device", type=int, default=0)
      parser.add_argument("--log_steps", type=int, default=1)
      parser.add_argument("--num_layers", type=int, default=3)
      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=16 * 1024)
      parser.add_argument("--walk_length", type=int, default=3)
      parser.add_argument("--lr", type=float, default=0.001)
      parser.add_argument("--epochs", type=int, default=200)
      parser.add_argument("--num_steps", type=int, default=100)
      parser.add_argument("--eval_steps", type=int, default=10)
      parser.add_argument("--runs", type=int, default=10)
      args = parser.parse_args()
      print(args)
  
      device = (
          f"cuda:{args.device}" if torch.cuda.is_available() else "cpu"
      )
      device = torch.device(device)
  
      dataset = PygLinkPropPredDataset(name="ogbl-citation2")
      split_edge = dataset.get_edge_split()
      data = dataset[0]
      data.edge_index = to_undirected(data.edge_index, data.num_nodes)
  
      loader = GraphSAINTRandomWalkSampler(
          data,
          batch_size=args.batch_size,
          walk_length=args.walk_length,
          num_steps=args.num_steps,
          sample_coverage=0,
          save_dir=dataset.processed_dir,
      )
  
      # We randomly pick some training samples that we want to evaluate on:
      torch.manual_seed(12345)
      idx = torch.randperm(split_edge["train"]["source_node"].numel())[
          :86596
      ]
      split_edge["eval_train"] = {
          "source_node": split_edge["train"]["source_node"][idx],
          "target_node": split_edge["train"]["target_node"][idx],
          "target_node_neg": split_edge["valid"]["target_node_neg"],
      }
  
      model = GCN(
          data.x.size(-1),
          args.hidden_channels,
          args.hidden_channels,
          args.num_layers,
          args.dropout,
      ).to(device)
      predictor = LinkPredictor(
          args.hidden_channels,
          args.hidden_channels,
          1,
          args.num_layers,
          args.dropout,
      ).to(device)
  
      evaluator = Evaluator(name="ogbl-citation2")
      logger = Logger(args.runs, args)
  
      run_idx = 0
  
      while run_idx < args.runs:
          model.reset_parameters()
          predictor.reset_parameters()
          optimizer = torch.optim.Adam(
              list(model.parameters()) + list(predictor.parameters()),
              lr=args.lr,
          )
  
          run_success = True
          for epoch in range(1, 1 + args.epochs):
              loss = train(model, predictor, loader, optimizer, device)
              print(
                  f"Run: {run_idx + 1:02d}, Epoch: {epoch:02d}, Loss: {loss:.4f}"
              )
              if loss > 2.0:
                  run_success = False
                  logger.reset(run_idx)
                  print("Learning failed. Rerun...")
                  break
  
              if epoch > 49 and epoch % args.eval_steps == 0:
                  result = test(
                      model,
                      predictor,
                      data,
                      split_edge,
                      evaluator,
                      batch_size=64 * 1024,
                      device=device,
                  )
                  logger.add_result(run_idx, result)
  
                  train_mrr, valid_mrr, test_mrr = result
                  print(
                      f"Run: {run_idx + 1:02d}, "
                      f"Epoch: {epoch:02d}, "
                      f"Loss: {loss:.4f}, "
                      f"Train: {train_mrr:.4f}, "
                      f"Valid: {valid_mrr:.4f}, "
                      f"Test: {test_mrr:.4f}"
                  )
  
          print("GraphSAINT")
          if run_success:
              logger.print_statistics(run_idx)
              run_idx += 1
  
      print("GraphSAINT")
      logger.print_statistics()
  
imports: import argparse
  import torch
  import numpy
  import torch.nn.functional
  from torch.utils.data import DataLoader
  from torch_sparse import SparseTensor
  from torch_geometric.data import GraphSAINTRandomWalkSampler
  from torch_geometric.nn import GCNConv
  from torch_geometric.utils import to_undirected
  from ogb.linkproppred import PygLinkPropPredDataset, Evaluator
  from logger import Logger
label_desc: Use image node_id 3 for calling a global function with example usage: main() without return types