code: def train(
      net, trainloader, validloader, optimizer, criterion, epoch, device
  ):
      best_model = pickle.dumps(net.state_dict())
  
      best_acc = 0.0
      for e in range(epoch):
          net.train()
          for graphs, labels in trainloader:
              labels = labels.to(device)
              graphs = graphs.to(device)
              feat = graphs.ndata.pop("attr")
              outputs = net(graphs, feat)
  
              loss = criterion(outputs, labels)
  
              # backprop
              optimizer.zero_grad()
              loss.backward()
              optimizer.step()
  
          gt = []
          pr = []
          net.eval()
          for graphs, labels in validloader:
              labels = labels.to(device)
              graphs = graphs.to(device)
              gt.append(labels)
              feat = graphs.ndata.pop("attr")
              outputs = net(graphs, feat)
              pr.append(outputs.argmax(1))
          gt = torch.cat(gt, dim=0)
          pr = torch.cat(pr, dim=0)
          acc = (gt == pr).float().mean().item()
          if acc > best_acc:
              best_acc = acc
              best_model = pickle.dumps(net.state_dict())
  
      net.load_state_dict(pickle.loads(best_model))
  
      return net
  
imports: import pickle
  from dgl.dataloading import GraphDataLoader
  import numpy
  from tqdm import tqdm
  import random
  import torch
  import torch.nn
  import torch.optim
  from dgl.data import GINDataset
  import torch
  import torch.nn
  import torch.nn.functional
  from dgl.nn.pytorch.conv import GINConv
  from dgl.nn.pytorch.glob import SumPooling, AvgPooling, MaxPooling
label_desc: Use image node_id 2 for calling a global function with example usage: train(net, trainloader, validloader, optimizer, criterion, epoch, device) and returns: net