code: def train(epoch):
      model.train()
  
      pbar = tqdm(total=paper_train_idx.size(0))
      pbar.set_description(f"Epoch {epoch:02d}")
  
      total_loss = 0
      for batch_size, n_id, adjs in train_loader:
          n_id = n_id.to(device)
          adjs = [adj.to(device) for adj in adjs]
          optimizer.zero_grad()
          out = model(
              n_id, x_dict, adjs, edge_type, node_type, local_node_idx
          )
          y = y_global[n_id][:batch_size].squeeze()
          loss = F.nll_loss(out, y)
          loss.backward()
          optimizer.step()
  
          total_loss += loss.item() * batch_size
          pbar.update(batch_size)
  
      pbar.close()
  
      loss = total_loss / paper_train_idx.size(0)
  
      return loss
  
imports: from copy import copy
  import argparse
  from tqdm.auto import tqdm
  import torch
  import torch.nn.functional
  from torch.nn import ModuleList, Linear, ParameterDict, Parameter
  from torch_sparse import SparseTensor
  from torch_geometric.utils import to_undirected
  from torch_geometric.data import NeighborSampler
  from torch_geometric.utils.hetero import group_hetero_graph
  from torch_geometric.nn import MessagePassing
  from ogb.nodeproppred import PygNodePropPredDataset, Evaluator
  from logger import Logger
label_desc: Use image node_id 1 for calling a global function with example usage: train(epoch) and returns: loss