code: def get_dataset_with_transformation(
      name, root, use_sparse_tensor=False, bf16=False
  ):
      path = osp.join(osp.dirname(osp.realpath(__file__)), root, name)
      transform = (
          T.ToSparseTensor(remove_edge_index=False)
          if use_sparse_tensor
          else None
      )
      if name == "ogbn-mag":
          if transform is None:
              transform = T.ToUndirected(merge=True)
          else:
              transform = T.Compose(
                  [T.ToUndirected(merge=True), transform]
              )
          dataset = OGB_MAG(
              root=path, preprocess="metapath2vec", transform=transform
          )
      elif name == "ogbn-products":
          if transform is None:
              transform = T.RemoveDuplicatedEdges()
          else:
              transform = T.Compose(
                  [T.RemoveDuplicatedEdges(), transform]
              )
  
          dataset = PygNodePropPredDataset(
              "ogbn-products", root=path, transform=transform
          )
  
      elif name == "Reddit":
          dataset = Reddit(root=path, transform=transform)
  
      data = dataset[0]
  
      if name == "ogbn-products":
          split_idx = dataset.get_idx_split()
          data.train_mask = index_to_mask(
              split_idx["train"], size=data.num_nodes
          )
          data.val_mask = index_to_mask(
              split_idx["valid"], size=data.num_nodes
          )
          data.test_mask = index_to_mask(
              split_idx["test"], size=data.num_nodes
          )
          data.y = data.y.squeeze()
  
      if bf16:
          if isinstance(data, HeteroData):
              for node_type in data.node_types:
                  data[node_type].x = data[node_type].x.to(
                      torch.bfloat16
                  )
          else:
              data.x = data.x.to(torch.bfloat16)
  
      return data, dataset.num_classes, transform
  
file_path: pytorch_geometric/benchmark/utils/utils.py
imports: import os
  import os.path
  from datetime import datetime
  import torch
  from ogb.nodeproppred import PygNodePropPredDataset
  from tqdm import tqdm
  import torch_geometric.transforms
  from torch_geometric.data import HeteroData
  from torch_geometric.datasets import OGB_MAG, Reddit
  from torch_geometric.nn import GAT, GCN, PNA, EdgeCNN, GraphSAGE
  from torch_geometric.utils import index_to_mask
  from .hetero_gat import HeteroGAT
  from .hetero_sage import HeteroGraphSAGE
label_desc: Use image node_id 1 for calling a global function with example usage: get_dataset_with_transformation(name, root, use_sparse_tensor, bf16) and returns: data,  dataset,  transform