code: def train(rank, world_size, dataset_name, root):
      dist.init_process_group(
          "nccl",
          "tcp://127.0.0.1:12347",
          world_size=world_size,
          rank=rank,
      )
      torch.cuda.set_device(rank)
  
      dataset = AsGraphPredDataset(
          DglGraphPropPredDataset(dataset_name, root)
      )
      evaluator = Evaluator(dataset_name)
  
      model = GIN(300, dataset.num_tasks).to(rank)
      model = nn.parallel.DistributedDataParallel(
          model, device_ids=[rank]
      )
      optimizer = optim.Adam(model.parameters(), lr=0.001)
      scheduler = optim.lr_scheduler.StepLR(
          optimizer, step_size=50, gamma=0.5
      )
  
      train_dataloader = GraphDataLoader(
          dataset[dataset.train_idx],
          batch_size=256,
          use_ddp=True,
          shuffle=True,
      )
      valid_dataloader = GraphDataLoader(
          dataset[dataset.val_idx], batch_size=256
      )
      test_dataloader = GraphDataLoader(
          dataset[dataset.test_idx], batch_size=256
      )
  
      for epoch in range(50):
          model.train()
          train_dataloader.set_epoch(epoch)
          for batched_graph, labels in train_dataloader:
              batched_graph, labels = batched_graph.to(rank), labels.to(
                  rank
              )
              node_feat, edge_feat = (
                  batched_graph.ndata["feat"],
                  batched_graph.edata["feat"],
              )
              logits = model(batched_graph, node_feat, edge_feat)
              optimizer.zero_grad()
              is_labeled = labels == labels
              loss = F.binary_cross_entropy_with_logits(
                  logits.float()[is_labeled], labels.float()[is_labeled]
              )
              loss.backward()
              optimizer.step()
          scheduler.step()
  
          if rank == 0:
              val_metric = evaluate(
                  valid_dataloader, rank, model.module, evaluator
              )[evaluator.eval_metric]
              test_metric = evaluate(
                  test_dataloader, rank, model.module, evaluator
              )[evaluator.eval_metric]
  
              print(
                  f"Epoch: {epoch:03d}, Loss: {loss:.4f}, "
                  f"Val: {val_metric:.4f}, Test: {test_metric:.4f}"
              )
  
      dist.destroy_process_group()
  
imports: import argparse
  import dgl
  import dgl.nn
  import torch
  import torch.distributed
  import torch.nn
  import torch.nn.functional
  import torch.optim
  from dgl.data import AsGraphPredDataset
  from dgl.dataloading import GraphDataLoader
  from ogb.graphproppred import DglGraphPropPredDataset, Evaluator
  from ogb.graphproppred.mol_encoder import AtomEncoder, BondEncoder
  from tqdm import tqdm
label_desc: Use image node_id 2 for calling a global function with example usage: train(rank, world_size, dataset_name, root) without return types