code: def collate_fn(batch):
      graphs, pmpds, labels = zip(*batch)
      batched_graphs = dgl.batch(graphs)
      batched_pmpds = sp.block_diag(pmpds)
      batched_labels = np.concatenate(labels, axis=0)
      return batched_graphs, batched_pmpds, batched_labels
  
imports: import os
  import dgl
  import torch
  import torch
  import torch.nn
  import torch.nn.functional
  from dgl.data import citation_graph
  import matplotlib.pyplot
  import networkx
  import dgl.function
  from torch.utils.data import DataLoader
label_desc: Use image node_id 4 for calling a global function with example usage: collate_fn(batch) and returns: batched_graphs,  batched_pmpds,  batched_labels