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 _train_epoch(self, loader, optimizer):
      self.generator.train()
  
      for iter, data_batch in enumerate(loader):
          anchor_data = data_batch
          anchor_data = anchor_data.to(self.device)
          g_loss_sum, reward_avg = 0, 0
  
          for _ in range(self.conf[GENERATOR_STEPS]):
              optimizer.zero_grad()
  
              neg_data, log_likelihoods = self.generator(anchor_data)
              if hasattr(self, "post_trans"):
                  neg_data = self.post_trans(neg_data)
  
              neg_out = self.reward_generator(
                  anchor_data, neg_data
              ).view(-1)
              rewards = torch.log(neg_out)
              if self.conf[BASELINE] == BaselineType.MEAN:
                  baseline = torch.mean(rewards)
                  advantages = (rewards - baseline).detach()
              elif self.conf[BASELINE] == BaselineType.EXP:
                  new_mean = torch.mean(rewards)
                  if self.moving_mean is None:
                      self.moving_mean = new_mean.cpu().detach().item()
                  else:
                      new_mean = new_mean.cpu().detach().item()
                      self.moving_mean = (
                          self.conf[MOVING_RATIO] * self.moving_mean
                          + (1.0 - self.conf[MOVING_RATIO]) * new_mean
                      )
                  advantages = (rewards - self.moving_mean).detach()
              elif isinstance(self.conf[BASELINE], float):
                  advantages = (rewards - self.conf[BASELINE]).detach()
              else:
                  advantages = rewards.detach()
              g_loss = -torch.sum(log_likelihoods.view(-1) * advantages)
  
              if torch.isnan(g_loss) or torch.isinf(g_loss):
                  torch.save(
                      self.generator.state_dict(),
                      os.path.join(self.out_path, "buggy_model.pt"),
                  )
                  exit()
  
              g_loss.backward()
              optimizer.step()
              g_loss_sum += g_loss.detach().item()
              reward_avg += torch.mean(rewards).detach().item()
  
          print(
              "Generator loss {}, Average reward {}".format(
                  g_loss_sum, reward_avg / self.conf[GENERATOR_STEPS]
              )
          )
  
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 4 for calling the RunnerGenerator obj's underlying member method code with example usage: obj._train_epoch(loader, optimizer) without return types