class_docstr: Runs the training of a graph classification model using the augmented
  data generated by the already trained generator. Check
  :obj:`examples.auggraph.GraphAug.run_aug_cls` for examples on how to run
  this augmented classifier 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.aug_cls_conf` for examples on
          how to define the conf dictionary for the augmented classifier
          model.
class_docstr_tok: ["Runs","the","training","of","a","graph","classification","model","using","the","augmented","data","generated","by","the","already","trained","generator",".","Check",":","obj",":","`","examples.auggraph.GraphAug.run_aug_cls","`","for","examples","on","how","to","run","this","augmented","classifier","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.aug_cls_conf","`","for","examples","on","how","to","define","the","conf","dictionary","for","the","augmented","classifier","model","."]
code: def _get_dataset(self, data_root_path, dataset_name):
      dataset = TUDataset(data_root_path, name=dataset_name.value)
      if dataset_name in [DatasetName.MUTAG]:
          self.data_trans = None
          self.conf[IN_DIMENSION] = dataset[0].x.shape[1]
          self.conf[EDGE_IN_DIMENSION] = dataset[0].x.shape[1]
      if dataset_name in [
          DatasetName.NCI1,
          DatasetName.NCI109,
          DatasetName.PROTEINS,
      ]:
          self.data_trans = None
          self.conf[IN_DIMENSION] = dataset[0].x.shape[1]
      elif dataset_name in [
          DatasetName.COLLAB,
          DatasetName.IMDB_BINARY,
      ]:
          self.data_trans = DegreeTrans(dataset)
          self.conf[IN_DIMENSION] = self.data_trans(dataset[0]).x.shape[
              1
          ]
      self.conf[NUM_CLASSES] = dataset.num_classes
      return dataset
  
imports: import os
  import torch
  import torch.nn.functional
  import numpy
  from torch_geometric.loader import DataLoader
  from torch_geometric.datasets import TUDataset
  from sklearn.model_selection import KFold, train_test_split
  from dig.auggraph.method.GraphAug.model import GIN, GCN
  from dig.auggraph.dataset.aug_dataset import DegreeTrans, Subset, AUG_trans
  from dig.auggraph.method.GraphAug.aug import Augmenter
  from dig.auggraph.method.GraphAug.constants import *
label_desc: Use image node_id 2 for calling the RunnerAugCls obj's underlying member method code with example usage: obj._get_dataset(data_root_path, dataset_name) and returns: dataset