code: def main(args):
      if not torch.cuda.is_available():
          args.mode = "cpu-cpu"
      print(f"Training in {args.mode} mode.")
      args.storage_device, args.device = args.mode.split("-")
      args.device = torch.device(args.device)
  
      # Load and preprocess dataset.
      print("Loading data")
      dataset = gb.BuiltinDataset("ogbl-citation2").load()
  
      # Move the dataset to the selected storage.
      graph = dataset.graph.to(args.storage_device)
      features = dataset.feature.to(args.storage_device)
  
      train_set = dataset.tasks[0].train_set
      args.fanout = list(map(int, args.fanout.split(",")))
  
      in_size = features.size("node", None, "feat")[0]
      hidden_channels = 256
      args.device = torch.device(args.device)
      model = SAGE(in_size, hidden_channels).to(args.device)
  
      # Model training.
      print("Training...")
      train(args, model, graph, features, train_set)
  
      # Test the model.
      print("Testing...")
      test_set = dataset.tasks[0].test_set
      valid_set = dataset.tasks[0].validation_set
      all_nodes_set = dataset.all_nodes_set
      valid_mrr, test_mrr = evaluate(
          args,
          model,
          graph,
          features,
          all_nodes_set,
          valid_set,
          test_set,
      )
      print(
          f"Validation MRR {valid_mrr.item():.4f}, "
          f"Test MRR {test_mrr.item():.4f}"
      )
  
imports: import argparse
  import time
  from functools import partial
  import dgl.graphbolt
  import dgl.nn
  import torch
  import torch.nn
  import torch.nn.functional
  import tqdm
  from ogb.linkproppred import Evaluator
label_desc: Use image node_id 6 for calling a global function with example usage: main(args) without return types