code: def train(
      dataset, model, prog_args, same_feat=True, val_dataset=None
  ):
      """
      training function
      """
      dir = prog_args.save_dir + "/" + prog_args.dataset
      if not os.path.exists(dir):
          os.makedirs(dir)
      dataloader = dataset
      optimizer = torch.optim.Adam(
          filter(lambda p: p.requires_grad, model.parameters()),
          lr=0.001,
      )
      early_stopping_logger = {"best_epoch": -1, "val_acc": -1}
  
      if prog_args.cuda > 0:
          torch.cuda.set_device(0)
      for epoch in range(prog_args.epoch):
          begin_time = time.time()
          model.train()
          accum_correct = 0
          total = 0
          print("\nEPOCH ###### {} ######".format(epoch))
          computation_time = 0.0
          for batch_idx, (batch_graph, graph_labels) in enumerate(
              dataloader
          ):
              for key, value in batch_graph.ndata.items():
                  batch_graph.ndata[key] = value.float()
              graph_labels = graph_labels.long()
              if torch.cuda.is_available():
                  batch_graph = batch_graph.to(
                      torch.cuda.current_device()
                  )
                  graph_labels = graph_labels.cuda()
  
              model.zero_grad()
              compute_start = time.time()
              ypred = model(batch_graph)
              indi = torch.argmax(ypred, dim=1)
              correct = torch.sum(indi == graph_labels).item()
              accum_correct += correct
              total += graph_labels.size()[0]
              loss = model.loss(ypred, graph_labels)
              loss.backward()
              batch_compute_time = time.time() - compute_start
              computation_time += batch_compute_time
              nn.utils.clip_grad_norm_(
                  model.parameters(), prog_args.clip
              )
              optimizer.step()
  
          train_accu = accum_correct / total
          print(
              "train accuracy for this epoch {} is {:.2f}%".format(
                  epoch, train_accu * 100
              )
          )
          elapsed_time = time.time() - begin_time
          print(
              "loss {:.4f} with epoch time {:.4f} s & computation time {:.4f} s ".format(
                  loss.item(), elapsed_time, computation_time
              )
          )
          global_train_time_per_epoch.append(elapsed_time)
          if val_dataset is not None:
              result = evaluate(val_dataset, model, prog_args)
              print("validation  accuracy {:.2f}%".format(result * 100))
              if (
                  result >= early_stopping_logger["val_acc"]
                  and result <= train_accu
              ):
                  early_stopping_logger.update(
                      best_epoch=epoch, val_acc=result
                  )
                  if prog_args.save_dir is not None:
                      torch.save(
                          model.state_dict(),
                          prog_args.save_dir
                          + "/"
                          + prog_args.dataset
                          + "/model.iter-"
                          + str(early_stopping_logger["best_epoch"]),
                      )
              print(
                  "best epoch is EPOCH {}, val_acc is {:.2f}%".format(
                      early_stopping_logger["best_epoch"],
                      early_stopping_logger["val_acc"] * 100,
                  )
              )
          torch.cuda.empty_cache()
      return early_stopping_logger
  
imports: import argparse
  import os
  import random
  import time
  import dgl
  import dgl.function
  import networkx
  import numpy
  import torch
  import torch.nn
  import torch.nn.functional
  import torch.utils.data
  from data_utils import pre_process
  from dgl import DGLGraph
  from dgl.data import tu
  from model.encoder import DiffPool
label_desc: Use image node_id 4 for calling a global function with example usage: train(dataset, model, prog_args, same_feat, val_dataset) and returns: early_stopping_logger