class_docstr: Runs the training of a reward generation model which will be able to
  distinguish between graphs with different labels. Check
  :obj:`examples.auggraph.GraphAug.run_reward_gen` for examples on how
  to run the reward generation model.
  
  Args:
      data_root_path (string): Directory where datasets should be saved.
      dataset_name (:class:`dig.auggraph.method.GraphAug.constants.enums.DatasetName`):
          Name of the graph dataset.
      conf (dict): Hyperparameters for the model. Check
          :obj:`examples.auggraph.GraphAug.conf.reward_gen_conf` for examples
          on how to define the conf dictionary for the reward generator.
code: def _get_dataset(self, data_root_path, dataset_name):
      dataset = TUDataset(data_root_path, name=dataset_name.value)
      if dataset_name in [
          DatasetName.NCI1,
          DatasetName.MUTAG,
          DatasetName.PROTEINS,
          DatasetName.NCI109,
      ]:
          self.train_set = TripleSet(dataset)
          self.val_set = TripleSet(dataset)
      elif dataset_name in [
          DatasetName.COLLAB,
          DatasetName.IMDB_BINARY,
      ]:
          self.train_set = TripleSet(
              dataset, transform=DegreeTrans(dataset)
          )
          self.val_set = TripleSet(
              dataset, transform=DegreeTrans(dataset)
          )
      self.conf[REWARD_GEN_PARAMS][IN_DIMENSION] = self.train_set[0][
          0
      ].x.shape[1]
  
file_path: DIG/dig/auggraph/method/GraphAug/runner_reward_gen.py
imports: import os
  import torch
  import torch.nn.functional
  from torch_geometric.loader import DataLoader
  from torch_geometric.datasets import TUDataset
  from .model import RewardGenModel
  from dig.auggraph.dataset.aug_dataset import DegreeTrans, TripleSet
  from dig.auggraph.method.GraphAug.constants import *
label_desc: Use image node_id 2 for calling the RunnerRewardGen obj's underlying member method code with example usage: obj._get_dataset(data_root_path, dataset_name) without return types