code: def train(g, target_idx, labels, num_classes, train_mask, model):
      # define train idx, loss function and optimizer
      train_idx = torch.nonzero(train_mask, as_tuple=False).squeeze()
      loss_fcn = nn.CrossEntropyLoss()
      optimizer = torch.optim.Adam(
          model.parameters(), lr=1e-2, weight_decay=5e-4
      )
  
      model.train()
      for epoch in range(50):
          logits = model(g)
          logits = logits[target_idx]
          loss = loss_fcn(logits[train_idx], labels[train_idx])
          optimizer.zero_grad()
          loss.backward()
          optimizer.step()
          acc = accuracy(
              logits[train_idx].argmax(dim=1),
              labels[train_idx],
              task="multiclass",
              num_classes=num_classes,
          ).item()
          print(
              "Epoch {:05d} | Loss {:.4f} | Train Accuracy {:.4f} ".format(
                  epoch, loss.item(), acc
              )
          )
  
file_path: dgl/examples/pytorch/rgcn/entity.py
imports: import argparse
  import dgl
  import torch
  import torch.nn
  import torch.nn.functional
  from dgl.data.rdf import AIFBDataset, AMDataset, BGSDataset, MUTAGDataset
  from dgl.nn.pytorch import RelGraphConv
  from torchmetrics.functional import accuracy
label_desc: Use image node_id 2 for calling a global function with example usage: train(g, target_idx, labels, num_classes, train_mask, model) without return types