code: def train(model, dataset, evaluator, device):
      train_dataloader = GraphDataLoader(
          dataset[dataset.train_idx],
          batch_size=256,
          shuffle=True,
          collate_fn=collate_dgl,
      )
      valid_dataloader = GraphDataLoader(
          dataset[dataset.val_idx],
          batch_size=256,
          collate_fn=collate_dgl,
      )
      test_dataloader = GraphDataLoader(
          dataset[dataset.test_idx],
          batch_size=256,
          collate_fn=collate_dgl,
      )
      optimizer = optim.Adam(model.parameters(), lr=0.001)
      num_epochs = 50
      scheduler = optim.lr_scheduler.StepLR(
          optimizer, step_size=num_epochs, gamma=0.5
      )
      loss_fcn = nn.BCEWithLogitsLoss()
  
      for epoch in range(num_epochs):
          model.train()
          total_loss = 0.0
          for batched_g, labels in train_dataloader:
              batched_g, labels = batched_g.to(device), labels.to(
                  device
              )
              logits = model(
                  batched_g,
                  batched_g.ndata["feat"],
                  batched_g.ndata["PE"],
              )
              loss = loss_fcn(logits, labels.float())
              total_loss += loss.item()
              optimizer.zero_grad()
              loss.backward()
              optimizer.step()
          scheduler.step()
          avg_loss = total_loss / len(train_dataloader)
          val_metric = evaluate(
              model, valid_dataloader, evaluator, device
          )
          test_metric = evaluate(
              model, test_dataloader, evaluator, device
          )
          print(
              f"Epoch: {epoch:03d}, Loss: {avg_loss:.4f}, "
              f"Val: {val_metric:.4f}, Test: {test_metric:.4f}"
          )
  
imports: import dgl
  import dgl.nn
  import dgl.sparse
  import torch
  import torch.nn
  import torch.nn.functional
  import torch.optim
  from dgl.data import AsGraphPredDataset
  from dgl.dataloading import GraphDataLoader
  from ogb.graphproppred import collate_dgl, DglGraphPropPredDataset, Evaluator
  from ogb.graphproppred.mol_encoder import AtomEncoder
  from tqdm import tqdm
label_desc: Use image node_id 2 for calling a global function with example usage: train(model, dataset, evaluator, device) without return types