code: def prepare_data(dataset, prog_args, train=False, pre_process=None):
      """
      preprocess TU dataset according to DiffPool's paper setting and load dataset into dataloader
      """
      if train:
          shuffle = True
      else:
          shuffle = False
  
      if pre_process:
          pre_process(dataset, prog_args)
  
      # dataset.set_fold(fold)
      return dgl.dataloading.GraphDataLoader(
          dataset,
          batch_size=prog_args.batch_size,
          shuffle=shuffle,
          num_workers=prog_args.n_worker,
      )
  
imports: import argparse
  import os
  import random
  import time
  import dgl
  import dgl.function
  import networkx
  import numpy
  import torch
  import torch.nn
  import torch.nn.functional
  import torch.utils.data
  from data_utils import pre_process
  from dgl import DGLGraph
  from dgl.data import tu
  from model.encoder import DiffPool
label_desc: Use image node_id 2 for calling a global function with example usage: prepare_data(dataset, prog_args, train, pre_process) and returns: dgl