code: def train(args, dataset, g, feats, paper_offset):
      print("Loading masks and labels")
      train_idx = (
          torch.LongTensor(dataset.get_idx_split("train"))
          + paper_offset
      )
      valid_idx = (
          torch.LongTensor(dataset.get_idx_split("valid"))
          + paper_offset
      )
      label = dataset.paper_label
  
      print("Initializing dataloader...")
      sampler = dgl.dataloading.MultiLayerNeighborSampler([15, 25])
      train_collator = ExternalNodeCollator(
          g, train_idx, sampler, paper_offset, feats, label
      )
      valid_collator = ExternalNodeCollator(
          g, valid_idx, sampler, paper_offset, feats, label
      )
      train_dataloader = torch.utils.data.DataLoader(
          train_collator.dataset,
          batch_size=1024,
          shuffle=True,
          drop_last=False,
          collate_fn=train_collator.collate,
          num_workers=4,
      )
      valid_dataloader = torch.utils.data.DataLoader(
          valid_collator.dataset,
          batch_size=1024,
          shuffle=True,
          drop_last=False,
          collate_fn=valid_collator.collate,
          num_workers=2,
      )
  
      print("Initializing model...")
      model = RGAT(
          dataset.num_paper_features,
          dataset.num_classes,
          1024,
          5,
          2,
          4,
          0.5,
          "paper",
      ).cuda()
      opt = torch.optim.Adam(model.parameters(), lr=0.001)
      sched = torch.optim.lr_scheduler.StepLR(
          opt, step_size=25, gamma=0.25
      )
  
      best_acc = 0
  
      for _ in range(args.epochs):
          model.train()
          with tqdm.tqdm(train_dataloader) as tq:
              for i, (input_nodes, output_nodes, mfgs) in enumerate(tq):
                  mfgs = [g.to("cuda") for g in mfgs]
                  x = mfgs[0].srcdata["x"]
                  y = mfgs[-1].dstdata["y"]
                  y_hat = model(mfgs, x)
                  loss = F.cross_entropy(y_hat, y)
                  opt.zero_grad()
                  loss.backward()
                  opt.step()
                  acc = (y_hat.argmax(1) == y).float().mean()
                  tq.set_postfix(
                      {
                          "loss": "%.4f" % loss.item(),
                          "acc": "%.4f" % acc.item(),
                      },
                      refresh=False,
                  )
  
          model.eval()
          correct = total = 0
          for i, (input_nodes, output_nodes, mfgs) in enumerate(
              tqdm.tqdm(valid_dataloader)
          ):
              with torch.no_grad():
                  mfgs = [g.to("cuda") for g in mfgs]
                  x = mfgs[0].srcdata["x"]
                  y = mfgs[-1].dstdata["y"]
                  y_hat = model(mfgs, x)
                  correct += (y_hat.argmax(1) == y).sum().item()
                  total += y_hat.shape[0]
          acc = correct / total
          print("Validation accuracy:", acc)
  
          sched.step()
  
          if best_acc < acc:
              best_acc = acc
              print("Updating best model...")
              torch.save(model.state_dict(), args.model_path)
  
imports: import argparse
  import time
  import dgl
  import dgl.function
  import dgl.nn
  import numpy
  import ogb
  import torch
  import torch.nn
  import torch.nn.functional
  import tqdm
  from ogb.lsc import MAG240MDataset, MAG240MEvaluator
label_desc: Use image node_id 1 for calling a global function with example usage: train(args, dataset, g, feats, paper_offset) without return types