code: def main(args):
      # Step 1: Prepare graph data and retrieve train/validation/test index ============================= #
      # Load dataset
      dataset = GASDataset(args.dataset)
      graph = dataset[0]
  
      # generate mini-batch only for forward edges
      sampler = dgl.dataloading.MultiLayerNeighborSampler([10, 10])
      tr_eid_dict = {}
      val_eid_dict = {}
      test_eid_dict = {}
      tr_eid_dict["forward"] = (
          graph.edges["forward"].data["train_mask"].nonzero().squeeze()
      )
      val_eid_dict["forward"] = (
          graph.edges["forward"].data["val_mask"].nonzero().squeeze()
      )
      test_eid_dict["forward"] = (
          graph.edges["forward"].data["test_mask"].nonzero().squeeze()
      )
  
      sampler = dgl.dataloading.as_edge_prediction_sampler(sampler)
      tr_loader = dgl.dataloading.DataLoader(
          graph,
          tr_eid_dict,
          sampler,
          batch_size=args.batch_size,
          shuffle=True,
          drop_last=False,
          num_workers=args.num_workers,
      )
      val_loader = dgl.dataloading.DataLoader(
          graph,
          val_eid_dict,
          sampler,
          batch_size=args.batch_size,
          shuffle=True,
          drop_last=False,
          num_workers=args.num_workers,
      )
      test_loader = dgl.dataloading.DataLoader(
          graph,
          test_eid_dict,
          sampler,
          batch_size=args.batch_size,
          shuffle=True,
          drop_last=False,
          num_workers=args.num_workers,
      )
  
      # check cuda
      if args.gpu >= 0 and th.cuda.is_available():
          device = "cuda:{}".format(args.gpu)
      else:
          device = "cpu"
  
      # binary classification
      num_classes = dataset.num_classes
  
      # Extract node features
      e_feats = graph.edges["forward"].data["feat"].shape[-1]
      u_feats = graph.nodes["u"].data["feat"].shape[-1]
      v_feats = graph.nodes["v"].data["feat"].shape[-1]
  
      # Step 2: Create model =================================================================== #
      model = GAS(
          e_in_dim=e_feats,
          u_in_dim=u_feats,
          v_in_dim=v_feats,
          e_hid_dim=args.e_hid_dim,
          u_hid_dim=args.u_hid_dim,
          v_hid_dim=args.v_hid_dim,
          out_dim=num_classes,
          num_layers=args.num_layers,
          dropout=args.dropout,
          activation=F.relu,
      )
  
      model = model.to(device)
  
      # Step 3: Create training components ===================================================== #
      loss_fn = th.nn.CrossEntropyLoss()
      optimizer = optim.Adam(
          model.parameters(), lr=args.lr, weight_decay=args.weight_decay
      )
  
      # Step 4: training epochs =============================================================== #
      for epoch in range(args.max_epoch):
          model.train()
          tr_loss = 0
          tr_f1 = 0
          tr_auc = 0
          tr_rap = 0
          tr_blocks = 0
          for input_nodes, edge_subgraph, blocks in tr_loader:
              blocks = [b.to(device) for b in blocks]
              edge_subgraph = edge_subgraph.to(device)
              u_feat = blocks[0].srcdata["feat"]["u"]
              v_feat = blocks[0].srcdata["feat"]["v"]
              f_feat = blocks[0].edges["forward"].data["feat"]
              b_feat = blocks[0].edges["backward"].data["feat"]
              labels = (
                  edge_subgraph.edges["forward"].data["label"].long()
              )
              logits = model(
                  edge_subgraph, blocks, f_feat, b_feat, u_feat, v_feat
              )
  
              # compute loss
              batch_loss = loss_fn(logits, labels)
              tr_loss += batch_loss.item()
              tr_f1 += f1_score(
                  labels.cpu(), logits.argmax(dim=1).cpu()
              )
              tr_auc += roc_auc_score(
                  labels.cpu(), logits[:, 1].detach().cpu()
              )
              tr_pre, tr_re, _ = precision_recall_curve(
                  labels.cpu(), logits[:, 1].detach().cpu()
              )
              tr_rap += tr_re[tr_pre > args.precision].max()
              tr_blocks += 1
  
              # backward
              optimizer.zero_grad()
              batch_loss.backward()
              optimizer.step()
  
          # validation
          model.eval()
          val_rap, val_f1, val_auc, val_loss = evaluate(
              model, loss_fn, val_loader, device
          )
  
          # Print out performance
          print(
              "In epoch {}, Train R@P: {:.4f} | Train F1: {:.4f} | Train AUC: {:.4f} | Train Loss: {:.4f}; "
              "Valid R@P: {:.4f} | Valid F1: {:.4f} | Valid AUC: {:.4f} | Valid loss: {:.4f}".format(
                  epoch,
                  tr_rap / tr_blocks,
                  tr_f1 / tr_blocks,
                  tr_auc / tr_blocks,
                  tr_loss / tr_blocks,
                  val_rap,
                  val_f1,
                  val_auc,
                  val_loss,
              )
          )
  
      # Test with mini batch after all epoch
      model.eval()
      test_rap, test_f1, test_auc, test_loss = evaluate(
          model, loss_fn, test_loader, device
      )
      print(
          "Test R@P: {:.4f} | Test F1: {:.4f} | Test AUC: {:.4f} | Test loss: {:.4f}".format(
              test_rap, test_f1, test_auc, test_loss
          )
      )
  
file_path: dgl/examples/pytorch/gas/main_sampling.py
imports: import argparse
  import dgl
  import torch
  import torch.nn.functional
  import torch.optim
  from dataloader import GASDataset
  from model_sampling import GAS
  from sklearn.metrics import f1_score, precision_recall_curve, roc_auc_score
label_desc: Use image node_id 2 for calling a global function with example usage: main(args) without return types