code: def train_model(
      model, g, opt, train_dataloader, num_epochs, rank, val_nid
  ):
      g.ndata["feat"]["_N"] = g.ndata["feat"]["_N"].to("cuda")
      g.ndata["label"]["_N"] = g.ndata["label"]["_N"].to("cuda")
      st = time.time()
      model.train()
      for epoch in range(num_epochs):
          total_loss = 0
          for _, (input_nodes, output_nodes, blocks) in enumerate(
              train_dataloader
          ):
              x = g.ndata["feat"]["_N"][input_nodes]
              y = g.ndata["label"]["_N"][output_nodes]
              y_hat = model(blocks, x)
              y = y.squeeze(1)
              loss = F.cross_entropy(y_hat, y)
              opt.zero_grad()
              loss.backward()
              opt.step()
              total_loss += loss.item()
          print(
              f"total loss: {total_loss} for epoch = {epoch} for rank = {rank}",
              flush=True,
          )
      et = time.time()
      print(
          f"Total time taken for num_epochs {num_epochs} "
          f"with batch_size {train_dataloader._batch_size} = {et-st} s on rank ={rank}"
      )
      if rank == 0:
          val_acc = layerwise_infer(g, val_nid, model, 1024 * 5, "cuda")
          print("---" * 30)
          print("Validation Accuracy {:.4f}".format(val_acc))
  
file_path: cugraph/python/cugraph-dgl/examples/multi_trainer_MG_example/model.py
imports: import time
  import dgl
  import torch
  import torch.nn
  import torch.nn.functional
  import torchmetrics.functional
  from cugraph_dgl.nn import SAGEConv
  import tqdm
label_desc: Use image node_id 2 for calling a global function with example usage: train_model(model, g, opt, train_dataloader, num_epochs, rank, val_nid) without return types