code: def main(args):
      # Step 1: Prepare graph data and retrieve train/validation/test index ============================= #
      # check cuda
      if args.gpu >= 0 and th.cuda.is_available():
          device = "cuda:{}".format(args.gpu)
      else:
          device = "cpu"
  
      # construct graph, split in/out edges and prepare train/validation/test data_loader
      data = Data(
          args.dataset,
          args.lbl_smooth,
          args.num_workers,
          args.batch_size,
      )
      data_iter = data.data_iter  # train/validation/test data_loader
      graph = data.g.to(device)
      num_rel = th.max(graph.edata["etype"]).item() + 1
  
      # Compute in/out edge norms and store in edata
      graph = in_out_norm(graph)
  
      # Step 2: Create model =================================================================== #
      compgcn_model = CompGCN_ConvE(
          num_bases=args.num_bases,
          num_rel=num_rel,
          num_ent=graph.num_nodes(),
          in_dim=args.init_dim,
          layer_size=args.layer_size,
          comp_fn=args.opn,
          batchnorm=True,
          dropout=args.dropout,
          layer_dropout=args.layer_dropout,
          num_filt=args.num_filt,
          hid_drop=args.hid_drop,
          feat_drop=args.feat_drop,
          ker_sz=args.ker_sz,
          k_w=args.k_w,
          k_h=args.k_h,
      )
      compgcn_model = compgcn_model.to(device)
  
      # Step 3: Create training components ===================================================== #
      loss_fn = th.nn.BCELoss()
      optimizer = optim.Adam(
          compgcn_model.parameters(), lr=args.lr, weight_decay=args.l2
      )
  
      # Step 4: training epoches =============================================================== #
      best_mrr = 0.0
      kill_cnt = 0
      for epoch in range(args.max_epochs):
          # Training and validation using a full graph
          compgcn_model.train()
          train_loss = []
          t0 = time()
          for step, batch in enumerate(data_iter["train"]):
              triple, label = batch[0].to(device), batch[1].to(device)
              sub, rel, obj, label = (
                  triple[:, 0],
                  triple[:, 1],
                  triple[:, 2],
                  label,
              )
              logits = compgcn_model(graph, sub, rel)
  
              # compute loss
              tr_loss = loss_fn(logits, label)
              train_loss.append(tr_loss.item())
  
              # backward
              optimizer.zero_grad()
              tr_loss.backward()
              optimizer.step()
  
          train_loss = np.sum(train_loss)
  
          t1 = time()
          val_results = evaluate(
              compgcn_model, graph, device, data_iter, split="valid"
          )
          t2 = time()
  
          # validate
          if val_results["mrr"] > best_mrr:
              best_mrr = val_results["mrr"]
              th.save(
                  compgcn_model.state_dict(),
                  "comp_link" + "_" + args.dataset,
              )
              kill_cnt = 0
              print("saving model...")
          else:
              kill_cnt += 1
              if kill_cnt > 100:
                  print("early stop.")
                  break
          print(
              "In epoch {}, Train Loss: {:.4f}, Valid MRR: {:.5}, Train time: {}, Valid time: {}".format(
                  epoch,
                  train_loss,
                  val_results["mrr"],
                  t1 - t0,
                  t2 - t1,
              )
          )
  
      # test use the best model
      compgcn_model.eval()
      compgcn_model.load_state_dict(
          th.load("comp_link" + "_" + args.dataset)
      )
      test_results = evaluate(
          compgcn_model, graph, device, data_iter, split="test"
      )
      print(
          "Test MRR: {:.5}\n, MR: {:.10}\n, H@10: {:.5}\n, H@3: {:.5}\n, H@1: {:.5}\n".format(
              test_results["mrr"],
              test_results["mr"],
              test_results["hits@10"],
              test_results["hits@3"],
              test_results["hits@1"],
          )
      )
  
file_path: dgl/examples/pytorch/compGCN/main.py
imports: import argparse
  from time import time
  import numpy
  import torch
  import torch.optim
  from data_loader import Data
  from models import CompGCN_ConvE
  from utils import in_out_norm
label_desc: Use image node_id 3 for calling a global function with example usage: main(args) without return types