code: def get_split_masks(data, dataset_name):
      if dataset_name == "ogbn-mag":
          train_mask = ("paper", data["paper"].train_mask)
          test_mask = ("paper", data["paper"].test_mask)
          val_mask = ("paper", data["paper"].val_mask)
      else:
          train_mask = data.train_mask
          val_mask = data.val_mask
          test_mask = data.test_mask
      return train_mask, val_mask, test_mask
  
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 4 for calling a global function with example usage: get_split_masks(data, dataset_name) and returns: train_mask,  val_mask,  test_mask