code: def train_epoch(model, optimizer, device, data_loader):
      model.train()
      epoch_loss = 0
      epoch_train_AP = 0
      list_scores = []
      list_labels = []
      for iter, (
          batch_graphs,
          batch_labels,
          batch_snorm_n,
      ) in enumerate(data_loader):
          batch_graphs = batch_graphs.to(device)
          batch_x = batch_graphs.ndata["feat"]  # num x feat
          batch_snorm_n = batch_snorm_n.to(device)
          batch_labels = batch_labels.to(device)
          optimizer.zero_grad()
  
          batch_scores = model(batch_graphs, batch_x, batch_snorm_n)
  
          loss = model.loss(batch_scores, batch_labels)
          loss.backward()
          optimizer.step()
          epoch_loss += loss.item()
          list_scores.append(batch_scores)
          list_labels.append(batch_labels)
  
      epoch_loss /= iter + 1
  
      evaluator = Evaluator(name="ogbg-molpcba")
      epoch_train_AP = evaluator.eval(
          {
              "y_pred": torch.cat(list_scores),
              "y_true": torch.cat(list_labels),
          }
      )["ap"]
  
      return epoch_loss, epoch_train_AP
  
imports: import argparse
  import random
  import dgl
  import numpy
  import torch
  import torch.nn
  import torch.nn.functional
  import torch.optim
  from dgl.dataloading import GraphDataLoader
  from ogb.graphproppred import Evaluator
  from ogb.graphproppred.mol_encoder import AtomEncoder
  from preprocessing import prepare_dataset
  from torch.utils.data import Dataset
  from tqdm import tqdm
label_desc: Use image node_id 5 for calling a global function with example usage: train_epoch(model, optimizer, device, data_loader) and returns: epoch_loss,  epoch_train_AP