code: def train(
      g,
      model,
      node_embed,
      optimizer,
      train_loader,
      split_idx,
      labels,
      logger,
      device,
      run,
  ):
      print("start training...")
      category = "paper"
  
      for epoch in range(3):
          num_train = split_idx["train"][category].shape[0]
          pbar = tqdm(total=num_train)
          pbar.set_description(f"Epoch {epoch:02d}")
          model.train()
  
          total_loss = 0
  
          for input_nodes, seeds, blocks in train_loader:
              blocks = [blk.to(device) for blk in blocks]
              seeds = seeds[
                  category
              ]  # we only predict the nodes with type "category"
              batch_size = seeds.shape[0]
              input_nodes_indexes = input_nodes["paper"].to(g.device)
              seeds = seeds.to(labels.device)
  
              emb = extract_embed(node_embed, input_nodes)
              # Add the batch's raw "paper" features
              emb.update(
                  {
                      "paper": g.ndata["feat"]["paper"][
                          input_nodes_indexes
                      ]
                  }
              )
  
              emb = {k: e.to(device) for k, e in emb.items()}
              lbl = labels[seeds].to(device)
  
              optimizer.zero_grad()
              logits = model(emb, blocks)[category]
  
              y_hat = logits.log_softmax(dim=-1)
              loss = F.nll_loss(y_hat, lbl)
              loss.backward()
              optimizer.step()
  
              total_loss += loss.item() * batch_size
              pbar.update(batch_size)
  
          pbar.close()
          loss = total_loss / num_train
  
          result = test(g, model, node_embed, labels, device, split_idx)
          logger.add_result(run, result)
          train_acc, valid_acc, test_acc = result
          print(
              f"Run: {run + 1:02d}, "
              f"Epoch: {epoch +1 :02d}, "
              f"Loss: {loss:.4f}, "
              f"Train: {100 * train_acc:.2f}%, "
              f"Valid: {100 * valid_acc:.2f}%, "
              f"Test: {100 * test_acc:.2f}%"
          )
  
      return logger
  
imports: import argparse
  import itertools
  import sys
  import dgl
  import dgl.nn
  import psutil
  import torch
  import torch.nn
  import torch.nn.functional
  from dgl import AddReverse, Compose, ToSimple
  from dgl.nn import HeteroEmbedding
  from ogb.nodeproppred import DglNodePropPredDataset, Evaluator
  from tqdm import tqdm
label_desc: Use image node_id 4 for calling a global function with example usage: train(g, model, node_embed, optimizer, train_loader, split_idx, labels, logger, device, run) and returns: logger