code: def run(args: argparse.ArgumentParser):
      for dataset_name in args.datasets:
          print(f"Dataset: {dataset_name}")
          root = osp.join(args.root, dataset_name)
          transform = (
              T.ToSparseTensor(remove_edge_index=False)
              if args.use_sparse_tensor
              else None
          )
          if dataset_name == "mag":
              transform = (
                  T.ToUndirected(merge=True)
                  if transform is None
                  else T.Compose(
                      [T.ToUndirected(merge=True), transform]
                  )
              )
              dataset = OGB_MAG(root=root, transform=transform)
              train_idx = ("paper", dataset[0]["paper"].train_mask)
              eval_idx = ("paper", None)
              neighbor_sizes = (
                  args.hetero_neighbor_sizes
                  if args.hetero_neighbor_sizes
                  else None
              )
          else:
              dataset = PygNodePropPredDataset(
                  f"ogbn-{dataset_name}", root
              )
              split_idx = dataset.get_idx_split()
              train_idx = split_idx["train"]
              eval_idx = None
              neighbor_sizes = (
                  args.homo_neighbor_sizes
                  if args.homo_neighbor_sizes
                  else None
              )
  
          data = dataset[0].to(args.device)
          average_times = []
          profile = torch_profile() if args.profile else nullcontext()
          # run dataloader iteration
          if neighbor_sizes is not None:
              for num_neighbors in neighbor_sizes:
                  print(
                      f"Training sampling with {num_neighbors} neighbors"
                  )
                  for batch_size in args.batch_sizes:
                      train_loader = NeighborLoader(
                          data,
                          num_neighbors=num_neighbors,
                          input_nodes=train_idx,
                          batch_size=batch_size,
                          shuffle=True,
                          num_workers=args.num_workers,
                          subgraph_type=args.subgraph_type,
                      )
                      cpu_affinity = (
                          train_loader.enable_cpu_affinity(
                              args.loader_cores
                          )
                          if args.cpu_affinity
                          else nullcontext()
                      )
                      runtimes = []
                      num_iterations = 0
                      with profile, cpu_affinity:
                          for runit in range(args.runs):
                              start = default_timer()
                              for batch in tqdm.tqdm(train_loader):
                                  num_iterations += 1
                              stop = default_timer()
                              runtimes.append(round(stop - start, 3))
                          average_time = round(
                              sum(runtimes) / args.runs, 3
                          )
                          print(
                              f"batch size={batch_size}, "
                              f"iterations={num_iterations}, "
                              f"runtimes={runtimes}, "
                              f"average runtime={average_time}"
                          )
                          average_times.append(average_time)
          eval_batch_sizes = (
              args.eval_batch_sizes if args.eval_batch_sizes else None
          )
          if eval_batch_sizes is not None:
              print("Evaluation sampling with all neighbors")
              for batch_size in eval_batch_sizes:
                  subgraph_loader = NeighborLoader(
                      data,
                      num_neighbors=[-1],
                      input_nodes=eval_idx,
                      batch_size=batch_size,
                      shuffle=False,
                      num_workers=args.num_workers,
                  )
                  cpu_affinity = (
                      subgraph_loader.enable_cpu_affinity(
                          args.loader_cores
                      )
                      if args.cpu_affinity
                      else nullcontext()
                  )
                  runtimes = []
                  num_iterations = 0
                  with profile, cpu_affinity:
                      for runit in range(args.runs):
                          start = default_timer()
                          for batch in tqdm.tqdm(subgraph_loader):
                              num_iterations += 1
                          stop = default_timer()
                          runtimes.append(round(stop - start, 3))
                      average_time = round(sum(runtimes) / args.runs, 3)
                      print(
                          f"batch size={batch_size}, "
                          f"iterations={num_iterations}, "
                          f"runtimes={runtimes}, "
                          f"average runtime={average_time}"
                      )
                      average_times.append(average_time)
          print(f"Total time averages: {average_times}")
  
file_path: pytorch_geometric/benchmark/loader/neighbor_loader.py
imports: import argparse
  import ast
  import os.path
  from contextlib import nullcontext
  from timeit import default_timer
  import tqdm
  from ogb.nodeproppred import PygNodePropPredDataset
  import torch_geometric.transforms
  from torch_geometric.datasets import OGB_MAG
  from torch_geometric.loader import NeighborLoader
  from torch_geometric.profile import torch_profile
label_desc: Use image node_id 1 for calling a global function with example usage: run(args) without return types