code: def train(model, data, args):
      optimizer = torch.optim.Adam(
          model.parameters(), lr=args.lr, weight_decay=args.weight_decay
      )
  
      parameters = model.state_dict()
      best_acc = 0.0
      for epoch in range(args.epoch):
          model.train()
          optimizer.zero_grad()
          if hasattr(model, "cls_forward"):
              output = model.cls_forward(data)
          else:
              output = model(data)
          loss = F.nll_loss(
              output[data.train_mask], data.y[data.train_mask]
          )
          loss.backward()
          optimizer.step()
  
          val_acc = test(model, data, data.val_mask)
          if val_acc > best_acc:
              best_acc = val_acc
              parameters = pickle.dumps(model.state_dict())
  
      model.load_state_dict(pickle.loads(parameters))
      return model
  
file_path: AutoGL/test/performance/node_classification/pyg/model_decouple.py
imports: import os
  import pickle
  import numpy
  from tqdm import tqdm
  import torch
  import torch.nn.functional
  from torch_geometric.datasets import Planetoid
  import torch_geometric.transforms
  from autogl.module.model.encoders import GCNEncoderMaintainer, GATEncoderMaintainer, SAGEEncoderMaintainer
  from autogl.module.model.decoders import LogSoftmaxDecoderMaintainer
  from autogl.solver.utils import set_seed
  from helper import get_encoder_decoder_hp
  import logging
label_desc: Use image node_id 2 for calling a global function with example usage: train(model, data, args) and returns: model