code: def main():
      parser = argparse.ArgumentParser(
          description="OGBN-Papers100M (Node2Vec)"
      )
      parser.add_argument("--device", type=int, default=0)
      parser.add_argument("--embedding_dim", type=int, default=128)
      parser.add_argument("--walk_length", type=int, default=20)
      parser.add_argument("--context_size", type=int, default=10)
      parser.add_argument("--walks_per_node", type=int, default=10)
      parser.add_argument("--batch_size", type=int, default=256)
      parser.add_argument("--lr", type=float, default=0.01)
      parser.add_argument("--epochs", type=int, default=2)
      parser.add_argument("--log_steps", type=int, default=10)
      parser.add_argument("--dropedge_rate", type=float, default=0.4)
      args = parser.parse_args()
  
      device = (
          f"cuda:{args.device}" if torch.cuda.is_available() else "cpu"
      )
      device = torch.device(device)
  
      dataset = PygNodePropPredDataset(name="ogbn-papers100M")
  
      split_idx = dataset.get_idx_split()
  
      data = dataset[0]
  
      print("Making the graph undirected.")
      ### Randomly drop some edges to avoid segmentation fault
      data.edge_index, _ = dropout_edge(
          data.edge_index, p=args.dropedge_rate
      )
      data.edge_index = to_undirected(data.edge_index, data.num_nodes)
      filename = "data_dict.pt"
  
      print(data)
  
      model = Node2Vec(
          data.edge_index,
          args.embedding_dim,
          args.walk_length,
          args.context_size,
          args.walks_per_node,
          sparse=True,
      ).to(device)
  
      loader = model.loader(
          batch_size=args.batch_size, shuffle=True, num_workers=4
      )
      optimizer = torch.optim.SparseAdam(
          list(model.parameters()), lr=args.lr
      )
  
      print("Saving data_dict before training...")
      save_data_dict(model, data, split_idx, save_file=filename)
  
      model.train()
      for epoch in range(1, args.epochs + 1):
          for i, (pos_rw, neg_rw) in tqdm(enumerate(loader)):
              optimizer.zero_grad()
              loss = model.loss(pos_rw.to(device), neg_rw.to(device))
              loss.backward()
              optimizer.step()
  
              if (i + 1) % args.log_steps == 0:
                  print(
                      f"Epoch: {epoch:02d}, Step: {i+1:03d}/{len(loader)}, "
                      f"Loss: {loss:.4f}"
                  )
  
              if (i + 1) % 1000 == 0:  # Save model every 1000 steps.
                  print("Saving data dict...")
                  save_data_dict(
                      model, data, split_idx, save_file=filename
                  )
  
          print("Saving data dict...")
          save_data_dict(model, data, split_idx, save_file=filename)
  
file_path: ogb/examples/nodeproppred/papers100M/node2vec.py
imports: import argparse
  import torch
  from torch_geometric.nn import Node2Vec
  from torch_geometric.utils import to_undirected, dropout_edge
  from tqdm.auto import tqdm
  from ogb.nodeproppred import PygNodePropPredDataset
label_desc: Use image node_id 2 for calling a global function with example usage: main() without return types