code: def train(self):
      import logging
  
      logger = logging.getLogger("PyGTrainer")
      logger.info("Entered train loop")
  
      total_loss = 0.0
      num_batches = 0
  
      time_forward = 0.0
      time_backward = 0.0
      time_loader = 0.0
      time_feature_transfer = 0.0
      start_time = time.perf_counter()
      end_time_backward = start_time
  
      for epoch in range(self.num_epochs):
          with td.algorithms.join.Join(
              [self.model], divide_by_initial_world_size=False
          ):
              self.model.train()
              for iter_i, data in enumerate(
                  self.get_loader(epoch=epoch, stage="train")
              ):
                  loader_time_iter = (
                      time.perf_counter() - end_time_backward
                  )
                  time_loader += loader_time_iter
  
                  time_feature_transfer_start = time.perf_counter()
  
                  num_sampled_nodes = sum(
                      [
                          torch.as_tensor(n)
                          for n in data.num_sampled_nodes_dict.values()
                      ]
                  )
                  num_sampled_edges = sum(
                      [
                          torch.as_tensor(e)
                          for e in data.num_sampled_edges_dict.values()
                      ]
                  )
  
                  # FIXME find a way to get around this and not have to call extend_tensor
                  num_layers = len(self.model.module.convs)
                  num_sampled_nodes = extend_tensor(
                      num_sampled_nodes, num_layers + 1
                  )
                  num_sampled_edges = extend_tensor(
                      num_sampled_edges, num_layers
                  )
  
                  data = data.to_homogeneous().cuda()
                  time_feature_transfer_end = time.perf_counter()
                  time_feature_transfer += (
                      time_feature_transfer_end
                      - time_feature_transfer_start
                  )
  
                  num_batches += 1
                  if iter_i % 20 == 1:
                      time_forward_iter = time_forward / num_batches
                      time_backward_iter = time_backward / num_batches
  
                      total_time_iter = (
                          time.perf_counter() - start_time
                      ) / num_batches
                      logger.info(f"epoch {epoch}, iteration {iter_i}")
                      logger.info(
                          f"num sampled nodes: {num_sampled_nodes}"
                      )
                      logger.info(
                          f"num sampled edges: {num_sampled_edges}"
                      )
                      logger.info(f"time forward: {time_forward_iter}")
                      logger.info(
                          f"time backward: {time_backward_iter}"
                      )
                      logger.info(f"loader time: {loader_time_iter}")
                      logger.info(
                          f"feature transfer time: {time_feature_transfer / num_batches}"
                      )
                      logger.info(f"total time: {total_time_iter}")
  
                  y_true = data.y
                  x = data.x.to(torch.float32)
  
                  start_time_forward = time.perf_counter()
                  edge_index = (
                      data.edge_index
                      if "edge_index" in data
                      else data.adj_t
                  )
  
                  self.optimizer.zero_grad()
                  y_pred = self.model(
                      x,
                      edge_index,
                      num_sampled_nodes,
                      num_sampled_edges,
                  )
  
                  end_time_forward = time.perf_counter()
                  time_forward += end_time_forward - start_time_forward
  
                  if y_pred.shape[0] > len(y_true):
                      raise ValueError(
                          f"illegal shape: {y_pred.shape}; {y_true.shape}"
                      )
  
                  y_true = y_true[: y_pred.shape[0]]
  
                  y_true = F.one_hot(
                      y_true.to(torch.int64),
                      num_classes=self.dataset.num_labels,
                  ).to(torch.float32)
  
                  if y_true.shape != y_pred.shape:
                      raise ValueError(
                          f"y_true shape was {y_true.shape} "
                          f"but y_pred shape was {y_pred.shape} "
                          f"in iteration {iter_i} "
                          f"on rank {y_pred.device.index}"
                      )
  
                  start_time_backward = time.perf_counter()
                  loss = F.cross_entropy(y_pred, y_true)
  
                  self.optimizer.zero_grad()
                  loss.backward()
                  self.optimizer.step()
                  total_loss += loss.item()
                  end_time_backward = time.perf_counter()
                  time_backward += (
                      end_time_backward - start_time_backward
                  )
  
          end_time = time.perf_counter()
  
          # test
          from torchmetrics import Accuracy
  
          acc = Accuracy(
              task="multiclass", num_classes=self.dataset.num_labels
          ).cuda()
  
          with td.algorithms.join.Join(
              [self.model], divide_by_initial_world_size=False
          ):
              self.model.eval()
              if self.rank == 0:
                  acc_sum = 0.0
                  with torch.no_grad():
                      for i, batch in enumerate(
                          self.get_loader(epoch=epoch, stage="test")
                      ):
                          num_sampled_nodes = sum(
                              [
                                  torch.as_tensor(n)
                                  for n in batch.num_sampled_nodes_dict.values()
                              ]
                          )
                          num_sampled_edges = sum(
                              [
                                  torch.as_tensor(e)
                                  for e in batch.num_sampled_edges_dict.values()
                              ]
                          )
                          batch_size = num_sampled_nodes[0]
  
                          batch = batch.to_homogeneous().cuda()
  
                          batch.y = batch.y.to(torch.long)
                          out = self.model.module(
                              batch.x,
                              batch.edge_index,
                              num_sampled_nodes,
                              num_sampled_edges,
                          )
                          acc_sum += acc(
                              out[:batch_size].softmax(dim=-1),
                              batch.y[:batch_size],
                          )
                  print(
                      f"Accuracy: {acc_sum/(i) * 100.0:.4f}%",
                  )
  
          td.barrier()
  
      with td.algorithms.join.Join(
          [self.model], divide_by_initial_world_size=False
      ):
          self.model.eval()
          if self.rank == 0:
              acc_sum = 0.0
              with torch.no_grad():
                  for i, batch in enumerate(
                      self.get_loader(epoch=epoch, stage="val")
                  ):
                      num_sampled_nodes = sum(
                          [
                              torch.as_tensor(n)
                              for n in batch.num_sampled_nodes_dict.values()
                          ]
                      )
                      num_sampled_edges = sum(
                          [
                              torch.as_tensor(e)
                              for e in batch.num_sampled_edges_dict.values()
                          ]
                      )
                      batch_size = num_sampled_nodes[0]
  
                      batch = batch.to_homogeneous().cuda()
  
                      batch.y = batch.y.to(torch.long)
                      out = self.model.module(
                          batch.x,
                          batch.edge_index,
                          num_sampled_nodes,
                          num_sampled_edges,
                      )
                      acc_sum += acc(
                          out[:batch_size].softmax(dim=-1),
                          batch.y[:batch_size],
                      )
              print(
                  f"Validation Accuracy: {acc_sum/(i) * 100.0:.4f}%",
              )
  
      stats = {
          "Accuracy": float(acc_sum / (i) * 100.0)
          if self.rank == 0
          else 0.0,
          "# Batches": num_batches,
          "Loader Time": time_loader,
          "Feature Transfer Time": time_feature_transfer,
          "Forward Time": time_forward,
          "Backward Time": time_backward,
      }
      return stats
  
imports: from trainers import Trainer
  from trainers import extend_tensor
  from datasets import OGBNPapers100MDataset
  from models.pyg import GraphSAGE
  import torch
  import numpy
  import torch.distributed
  from torch.nn.parallel import DistributedDataParallel
  import torch.nn.functional
  from torch_geometric.utils.sparse import index2ptr
  from torch_geometric.data import HeteroData
  from torch_geometric.loader import NeighborLoader
  import gc
  import os
  import time
label_desc: Use image node_id 1 for calling the PyGTrainer obj's underlying member method code with example usage: obj.train() and returns: stats