code: def split(data, batch):
      # i-th contains elements from slice[i] to slice[i+1]
      node_slice = torch.cumsum(torch.from_numpy(np.bincount(batch)), 0)
      node_slice = torch.cat([torch.tensor([0]), node_slice])
      row, _ = data.edge_index
      edge_slice = torch.cumsum(
          torch.from_numpy(np.bincount(batch[row])), 0
      )
      edge_slice = torch.cat([torch.tensor([0]), edge_slice])
  
      # Edge indices should start at zero for every graph.
      data.edge_index -= node_slice[batch[row]].unsqueeze(0)
      data.__num_nodes__ = np.bincount(batch).tolist()
  
      slices = dict()
      slices["x"] = node_slice
      slices["edge_index"] = edge_slice
      slices["y"] = torch.arange(0, batch[-1] + 2, dtype=torch.long)
      return data, slices
  
imports: import os
  import glob
  import json
  import torch
  import pickle
  import numpy
  import os.path
  from torch_geometric.data import Data, InMemoryDataset
  import traceback
label_desc: Use image node_id 2 for calling a global function with example usage: split(data, batch) and returns: data,  slices