code: def train(model, G, args, train_mask, labels):
      optimizer = torch.optim.AdamW(model.parameters())
      scheduler = torch.optim.lr_scheduler.OneCycleLR(
          optimizer, total_steps=args.n_epoch, max_lr=args.max_lr
      )
      for epoch in np.arange(args.n_epoch) + 1:
          model.train()
          logits = model(G)
          # The loss is computed only for labeled nodes.
          loss = F.cross_entropy(
              logits[train_mask], labels[train_mask].to(args.device)
          )
          optimizer.zero_grad()
          loss.backward()
          torch.nn.utils.clip_grad_norm_(model.parameters(), args.clip)
          optimizer.step()
          scheduler.step()
  
file_path: AutoGL/test/performance/heterogeneous/dgl/hgt_main.py
imports: import os
  import numpy
  from tqdm import tqdm
  import torch
  import torch.nn.functional
  from autogl.module.model.dgl import AutoHGT, AutoHeteroRGCN
  from autogl.solver.utils import set_seed
  import numpy
  import argparse
  from autogl.datasets import build_dataset_from_name
label_desc: Use image node_id 3 for calling a global function with example usage: train(model, G, args, train_mask, labels) without return types