code: def paint(batched_points):
      B, N = batched_points.shape
      colored = color_map[batched_points].squeeze(2)
      return colored
  
file_path: dgl/examples/pytorch/pointcloud/point_transformer/train_partseg.py
imports: import argparse
  import time
  from functools import partial
  import dgl
  import numpy
  import torch
  import torch.optim
  import tqdm
  from point_transformer import PartSegLoss, PointTransformerSeg
  from ShapeNet import ShapeNet
  from torch.utils.data import DataLoader
label_desc: Use image node_id 5 for calling a global function with example usage: paint(batched_points) and returns: colored