class_docstr: Runs the training of an augmented samples generator model which uses the
  already trained reward generation model. For a given graph, the model
  generates an augmented sample and a likelihood that this is a label
  invariant augmentation. This prediction is then evaluated by the reward
  generation model and a loss is computed based on these metrics. The loss
  is then minimized through training. Check
  :obj:`examples.auggraph.GraphAug.run_generator` for examples on how to
  run the generator 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.generator_conf` for
          examples on how to define the conf dictionary for the generator.
class_docstr_tok: ["Runs","the","training","of","an","augmented","samples","generator","model","which","uses","the","already","trained","reward","generation","model",".","For","a","given","graph",",","the","model","generates","an","augmented","sample","and","a","likelihood","that","this","is","a","label","invariant","augmentation",".","This","prediction","is","then","evaluated","by","the","reward","generation","model","and","a","loss","is","computed","based","on","these","metrics",".","The","loss","is","then","minimized","through","training",".","Check",":","obj",":","`","examples.auggraph.GraphAug.run_generator","`","for","examples","on","how","to","run","the","generator","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.generator_conf","`","for","examples","on","how","to","define","the","conf","dictionary","for","the","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.AIDS,
          DatasetName.NCI109,
          DatasetName.PROTEINS,
      ]:
          self.train_set = Subset(dataset)
          self.val_set = Subset(dataset)
      elif dataset_name in [
          DatasetName.COLLAB,
          DatasetName.IMDB_BINARY,
      ]:
          self.train_set = Subset(
              dataset, transform=DegreeTrans(dataset)
          )
          self.val_set = Subset(dataset, transform=DegreeTrans(dataset))
          self.post_trans = DegreeTrans(dataset)
      in_dim = self.val_set[0].x.shape[1]
      self.conf[REWARD_GEN_PARAMS][IN_DIMENSION] = in_dim
      self.conf[GENERATOR_PARAMS][IN_DIMENSION] = in_dim
      if (
          AugType.NODE_FM.value
          in self.conf[GENERATOR_PARAMS][AUG_TYPE_PARAMS]
      ):
          self.conf[GENERATOR_PARAMS][AUG_TYPE_PARAMS][
              AugType.NODE_FM.value
          ][NODE_FEAT_DIM] = in_dim
  
file_path: DIG/dig/auggraph/method/GraphAug/runner_generator.py
imports: import os
  import torch
  from torch_geometric.loader import DataLoader
  from torch_geometric.data import Batch
  from torch_geometric.datasets import TUDataset
  from .model import RewardGenModel
  from dig.auggraph.dataset.aug_dataset import Subset, DegreeTrans
  from .aug import Augmenter
  from dig.auggraph.method.GraphAug.constants import *
label_desc: Use image node_id 2 for calling the RunnerGenerator obj's underlying member method code with example usage: obj._get_dataset(data_root_path, dataset_name) without return types