code: def train(
      dataloader,
      test_g,
      test_nids,
      test_mask,
      triplets,
      device,
      model_state_file,
      model,
  ):
      optimizer = torch.optim.Adam(model.parameters(), lr=1e-2)
      best_mrr = 0
      for epoch, batch_data in enumerate(
          dataloader
      ):  # single graph batch
          model.train()
          g, train_nids, edges, labels = batch_data
          g = g.to(device)
          train_nids = train_nids.to(device)
          edges = edges.to(device)
          labels = labels.to(device)
  
          embed = model(g, train_nids)
          loss = model.get_loss(embed, edges, labels)
          optimizer.zero_grad()
          loss.backward()
          nn.utils.clip_grad_norm_(
              model.parameters(), max_norm=1.0
          )  # clip gradients
          optimizer.step()
          print(
              "Epoch {:04d} | Loss {:.4f} | Best MRR {:.4f}".format(
                  epoch, loss.item(), best_mrr
              )
          )
          if (epoch + 1) % 500 == 0:
              # perform validation on CPU because full graph is too large
              model = model.cpu()
              model.eval()
              embed = model(test_g, test_nids)
              mrr = calc_mrr(
                  embed,
                  model.w_relation,
                  test_mask,
                  triplets,
                  batch_size=500,
              )
              # save best model
              if best_mrr < mrr:
                  best_mrr = mrr
                  torch.save(
                      {
                          "state_dict": model.state_dict(),
                          "epoch": epoch,
                      },
                      model_state_file,
                  )
              model = model.to(device)
  
imports: import dgl
  import numpy
  import torch
  import torch.nn
  import torch.nn.functional
  import tqdm
  from dgl.data.knowledge_graph import FB15k237Dataset
  from dgl.dataloading import GraphDataLoader
  from dgl.nn.pytorch import RelGraphConv
label_desc: Use image node_id 5 for calling a global function with example usage: train(dataloader, test_g, test_nids, test_mask, triplets, device, model_state_file, model) without return types