code: def get_dataset(name, dataset_dir, use_sparse_tensor=False):
      transforms = []
      if use_sparse_tensor:
          transforms = [T.ToSparseTensor(remove_edge_index=False)]
  
      if name == "ogbn-mag":
          transforms = [T.ToUndirected(merge=True)] + transforms
          return OGB_MAG(
              root=dataset_dir,
              preprocess="metapath2vec",
              transform=T.Compose(transforms),
          )
  
      elif name == "ogbn-products":
          transforms = [T.RemoveDuplicatedEdges()] + transforms
          return PygNodePropPredDataset(
              "ogbn-products",
              root=dataset_dir,
              transform=T.Compose(transforms),
          )
  
      elif name == "Reddit":
          return Reddit(
              root=dataset_dir,
              transform=T.Compose(transforms),
          )
  
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 2 for calling a global function with example usage: get_dataset(name, dataset_dir, use_sparse_tensor) and returns: OGB_MAG, PygNodePropPredDataset, Reddit