code: def train(model, device, loader, optimizer):
      model.train()
      loss_accum = 0
  
      for step, (bg, labels) in enumerate(
          tqdm(loader, desc="Iteration")
      ):
          bg = bg.to(device)
          x = bg.ndata.pop("feat")
          edge_attr = bg.edata.pop("feat")
          labels = labels.to(device)
  
          pred = model(bg, x, edge_attr).view(
              -1,
          )
          optimizer.zero_grad()
          loss = reg_criterion(pred, labels)
          loss.backward()
          optimizer.step()
  
          loss_accum += loss.detach().cpu().item()
  
      return loss_accum / (step + 1)
  
file_path: dgl/examples/pytorch/ogb_lsc/PCQM4M/main.py
imports: import argparse
  import os
  import random
  import dgl
  import numpy
  import torch
  import torch.optim
  from gnn import GNN
  from ogb.lsc import DglPCQM4MDataset, PCQM4MEvaluator
  from torch.optim.lr_scheduler import StepLR
  from torch.utils.data import DataLoader
  from torch.utils.tensorboard import SummaryWriter
  from tqdm import tqdm
label_desc: Use image node_id 2 for calling a global function with example usage: train(model, device, loader, optimizer) and returns: unknown