code: def transform_before_split(dataset):
      """
      Dataset transformation before train/val/test split
      :param dataset: A DeepSNAP dataset object
      :return: A transformed DeepSNAP dataset object
      """
      if cfg.dataset.remove_feature:
          dataset.apply_transform(
              remove_node_feature,
              update_graph=True,
              update_tensor=False,
          )
      augmentation = preprocess.FeatureAugment()
      actual_feat_dims, actual_label_dim = augmentation.augment(dataset)
      if cfg.dataset.augment_label:
          dataset.apply_transform(
              preprocess._replace_label,
              update_graph=True,
              update_tensor=False,
          )
      # Update augmented feature/label dims by real dims (user specified dims
      # may not be realized)
      cfg.dataset.augment_feature_dims = actual_feat_dims
      if cfg.dataset.augment_label:
          cfg.dataset.augment_label_dims = actual_label_dim
  
      # Temporary for ID-GNN path prediction task
      if cfg.dataset.task == "edge" and "id" in cfg.gnn.layer_type:
          dataset.apply_transform(
              path_len, update_graph=False, update_tensor=False
          )
  
      return dataset
  
imports: import logging
  import pickle
  import time
  import networkx
  import torch
  import torch_geometric.transforms
  from deepsnap.batch import Batch
  from deepsnap.dataset import GraphDataset
  from ogb.graphproppred import PygGraphPropPredDataset
  from torch.utils.data import DataLoader
  from torch_geometric.datasets import PPI, Amazon, Coauthor, KarateClub, MNISTSuperpixels, Planetoid, QM7b, TUDataset
  import graphgym.models.feature_augment
  import graphgym.register
  from graphgym.config import cfg
  from graphgym.models.transform import edge_nets, ego_nets, path_len, remove_node_feature
label_desc: Use image node_id 5 for calling a global function with example usage: transform_before_split(dataset) and returns: dataset