code: def get_loader(self, epoch: int = 0, stage="train"):
      import logging
  
      logger = logging.getLogger("PyGNativeTrainer")
      logger.info(f"Getting loader for epoch {epoch}")
  
      if stage == "train":
          mask_dict = self.__dataset.train_dict
      elif stage == "test":
          mask_dict = self.__dataset.test_dict
      elif stage == "val":
          mask_dict = self.__dataset.val_dict
      else:
          raise ValueError(f"Invalid stage {stage}")
  
      input_nodes_dict = {
          node_type: np.array_split(
              np.arange(len(mask))[mask], self.__world_size
          )[self.__rank]
          for node_type, mask in mask_dict.items()
      }
  
      input_nodes = list(input_nodes_dict.items())
      if len(input_nodes) > 1:
          raise ValueError(
              "Multiple input node types currently unsupported"
          )
      else:
          input_nodes = tuple(input_nodes[0])
  
      # get loader
      loader = NeighborLoader(
          self.data,
          input_nodes=input_nodes,
          is_sorted=True,
          disjoint=False,
          num_workers=pyg_num_workers(
              self.__world_size
          ),  # FIXME change this
          persistent_workers=True,
          **self.__loader_kwargs,  # batch size, num neighbors, replace, shuffle, etc.
      )
  
      logger.info("done creating loader")
      return loader
  
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 8 for calling the PyGNativeTrainer obj's underlying member method code with example usage: obj.get_loader(epoch, stage) and returns: loader