code: def train(net, opt, scheduler, train_loader, dev):
      category_list = sorted(list(shapenet.seg_classes.keys()))
      eye_mat = np.eye(16)
      net.train()
  
      total_loss = 0
      num_batches = 0
      total_correct = 0
      count = 0
      start = time.time()
      with tqdm.tqdm(train_loader, ascii=True) as tq:
          for data, label, cat in tq:
              num_examples = data.shape[0]
              data = data.to(dev, dtype=torch.float)
              label = label.to(dev, dtype=torch.long).view(-1)
              opt.zero_grad()
              cat_ind = [category_list.index(c) for c in cat]
              # An one-hot encoding for the object category
              cat_tensor = (
                  torch.tensor(eye_mat[cat_ind])
                  .to(dev, dtype=torch.float)
                  .repeat(1, 2048)
              )
              cat_tensor = cat_tensor.view(num_examples, -1, 16)
              logits = net(data, cat_tensor).permute(0, 2, 1)
              loss = L(logits, label)
              loss.backward()
              opt.step()
  
              _, preds = logits.max(1)
  
              count += num_examples * 2048
              loss = loss.item()
              total_loss += loss
              num_batches += 1
              correct = (preds.view(-1) == label).sum().item()
              total_correct += correct
  
              AvgLoss = total_loss / num_batches
              AvgAcc = total_correct / count
  
              tq.set_postfix(
                  {
                      "AvgLoss": "%.5f" % AvgLoss,
                      "AvgAcc": "%.5f" % AvgAcc,
                  }
              )
      scheduler.step()
      end = time.time()
      print(
          "[Train] AvgLoss: {:.5}, AvgAcc: {:.5}, Time: {:.5}s".format(
              total_loss / num_batches,
              total_correct / count,
              end - start,
          )
      )
      return data, preds, AvgLoss, AvgAcc, end - start
  
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 2 for calling a global function with example usage: train(net, opt, scheduler, train_loader, dev) and returns: data,  preds,  AvgLoss,  AvgAcc,  unknown