code: def train(args, model, graph, features, train_set):
      optimizer = torch.optim.Adam(model.parameters(), lr=args.lr)
      dataloader = create_dataloader(args, graph, features, train_set)
  
      for epoch in range(args.epochs):
          model.train()
          total_loss = 0
          start_epoch_time = time.time()
          for step, data in tqdm.tqdm(enumerate(dataloader)):
              # Get node pairs with labels for loss calculation.
              compacted_pairs, labels = data.node_pairs_with_labels
  
              node_feature = data.node_features["feat"]
              blocks = data.blocks
  
              # Get the embeddings of the input nodes.
              y = model(blocks, node_feature)
              logits = model.predictor(
                  y[compacted_pairs[0]] * y[compacted_pairs[1]]
              ).squeeze()
  
              # Compute loss.
              loss = F.binary_cross_entropy_with_logits(logits, labels)
              optimizer.zero_grad()
              loss.backward()
              optimizer.step()
  
              total_loss += loss.item()
              if step + 1 == args.early_stop:
                  break
  
          end_epoch_time = time.time()
          print(
              f"Epoch {epoch:05d} | "
              f"Loss {(total_loss) / (step + 1):.4f} | "
              f"Time {(end_epoch_time - start_epoch_time):.4f} s"
          )
  
imports: import argparse
  import time
  from functools import partial
  import dgl.graphbolt
  import dgl.nn
  import torch
  import torch.nn
  import torch.nn.functional
  import tqdm
  from ogb.linkproppred import Evaluator
label_desc: Use image node_id 4 for calling a global function with example usage: train(args, model, graph, features, train_set) without return types