code: def train(model, graph, args, label, train_mask, val_mask):
      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()
          output = model(graph)
          loss = F.nll_loss(output[train_mask], label[train_mask])
          loss.backward()
          optimizer.step()
  
          val_acc = test(model, graph, val_mask, label)
          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/dgl/model.py
imports: import numpy
  from tqdm import tqdm
  import pickle
  import torch
  import torch.nn.functional
  from dgl.data import CoraGraphDataset, PubmedGraphDataset, CiteseerGraphDataset
  from autogl.module.model.dgl import AutoGCN, AutoGAT, AutoSAGE
  from autogl.solver.utils import set_seed
  import logging
label_desc: Use image node_id 2 for calling a global function with example usage: train(model, graph, args, label, train_mask, val_mask) and returns: model