code: def get_idx_split(dataset, dataset_name, split_data):
      if dataset_name == "ogbn-mag" or dataset_name == "Reddit":
          train_idx = mask_to_index(split_data.train_mask)
          test_idx = mask_to_index(split_data.test_mask)
          val_idx = mask_to_index(split_data.val_mask)
  
      elif dataset_name == "ogbn-products":
          split_idx = dataset.get_idx_split()
          train_idx = split_idx["train"]
          test_idx = split_idx["test"]
          val_idx = split_idx["valid"]
  
      return train_idx, val_idx, test_idx
  
file_path: pytorch_geometric/examples/distributed/pyg/partition_graph.py
imports: import argparse
  import os
  import os.path
  import torch
  from ogb.nodeproppred import PygNodePropPredDataset
  import torch_geometric.transforms
  from torch_geometric.datasets import OGB_MAG, Reddit
  from torch_geometric.distributed import Partitioner
  from torch_geometric.utils import mask_to_index
label_desc: Use image node_id 3 for calling a global function with example usage: get_idx_split(dataset, dataset_name, split_data) and returns: train_idx,  val_idx,  test_idx