code: def train(train_loader, val_loader, device, model):
      # loss function, optimizer and scheduler
      loss_fcn = nn.CrossEntropyLoss()
      optimizer = optim.Adam(model.parameters(), lr=0.01)
      scheduler = optim.lr_scheduler.StepLR(
          optimizer, step_size=50, gamma=0.5
      )
  
      # training loop
      for epoch in range(350):
          model.train()
          total_loss = 0
          for batch, (batched_graph, labels) in enumerate(train_loader):
              batched_graph = batched_graph.to(device)
              labels = labels.to(device)
              feat = batched_graph.ndata.pop("attr")
              logits = model(batched_graph, feat)
              loss = loss_fcn(logits, labels)
              optimizer.zero_grad()
              loss.backward()
              optimizer.step()
              total_loss += loss.item()
          scheduler.step()
          train_acc = evaluate(train_loader, device, model)
          valid_acc = evaluate(val_loader, device, model)
          print(
              "Epoch {:05d} | Loss {:.4f} | Train Acc. {:.4f} | Validation Acc. {:.4f} ".format(
                  epoch, total_loss / (batch + 1), train_acc, valid_acc
              )
          )
  
imports: import argparse
  import numpy
  import torch
  import torch.nn
  import torch.nn.functional
  import torch.optim
  from dgl.data import GINDataset
  from dgl.dataloading import GraphDataLoader
  from dgl.nn.pytorch.conv import GINConv
  from dgl.nn.pytorch.glob import SumPooling
  from sklearn.model_selection import StratifiedKFold
  from torch.utils.data.sampler import SubsetRandomSampler
label_desc: Use image node_id 3 for calling a global function with example usage: train(train_loader, val_loader, device, model) without return types