code: def train(model, train_loader, val_loader, args):
      optimizer = torch.optim.Adam(model.parameters(), lr=args.lr)
  
      parameters = model.state_dict()
      best_acc = 0.0
      for epoch in range(args.epoch):
          model.train()
          for data in train_loader:
              data = data.to(args.device)
              optimizer.zero_grad()
              output = model(data)
              loss = F.nll_loss(output, data.y)
              loss.backward()
              optimizer.step()
  
          val_acc = test(model, val_loader, args)
          if val_acc > best_acc:
              best_acc = val_acc
              parameters = model.state_dict()
  
      model.load_state_dict(parameters)
      return model
  
imports: import os
  import random
  import numpy
  from tqdm import tqdm
  import torch
  import torch.nn.functional
  from torch.nn import Sequential, Linear, ReLU
  import torch_geometric
  from torch_geometric.datasets import TUDataset
  from torch_geometric.nn import GINConv, global_add_pool, GraphConv, TopKPooling
  from torch_geometric.nn import global_mean_pool, global_max_pool
  import logging
label_desc: Use image node_id 2 for calling a global function with example usage: train(model, train_loader, val_loader, args) and returns: model