code: def main(args):
      if not torch.cuda.is_available():
          args.mode = "cpu"
      print(f"Training in {args.mode} mode.")
  
      # Load and preprocess dataset.
      print("Loading data")
      dataset = DglLinkPropPredDataset("ogbl-citation2")
      g = dataset[0]
      if args.compare_graphbolt:
          fused_sampling = False
      else:
          fused_sampling = True
          g = g.to("cuda" if args.mode == "puregpu" else "cpu")
  
      # Whether use Unified Virtual Addressing (UVA) for CUDA computation.
      use_uva = args.mode == "mixed"
      device = torch.device("cpu" if args.mode == "cpu" else "cuda")
  
      # Convert the graph to its bidirectional form.
      g, reverse_eids = to_bidirected_with_reverse_mapping(g)
      reverse_eids = reverse_eids.to(g.device)
      seed_edges = torch.arange(g.num_edges()).to(g.device)
      edge_split = dataset.get_edge_split()
  
      # Create GraphSAGE model.
      in_size = g.ndata["feat"].shape[1]
      model = SAGE(in_size, 256).to(device)
  
      # Model training.
      print("Training...")
      train(
          args,
          device,
          g,
          reverse_eids,
          seed_edges,
          model,
          use_uva,
          fused_sampling,
      )
  
      # Validate/Test the model.
      print("Validation/Testing...")
      valid_mrr, test_mrr = evaluate(
          device, g, edge_split, model, batch_size=args.eval_batch_size
      )
      print(
          f"Validation MRR {valid_mrr.item():.4f}, Test MRR {test_mrr.item():.4f}"
      )
  
imports: import argparse
  import time
  import dgl
  import dgl.nn
  import torch
  import torch.nn
  import torch.nn.functional
  import tqdm
  from dgl.dataloading import as_edge_prediction_sampler, DataLoader, MultiLayerFullNeighborSampler, negative_sampler, NeighborSampler
  from ogb.linkproppred import DglLinkPropPredDataset, Evaluator
label_desc: Use image node_id 6 for calling a global function with example usage: main(args) without return types