code: def __init__(
      self,
      in_dim,
      num_layers,
      hidden,
      pool_type=PoolType.SUM,
      model_type=RewardGenModelType.GMNET,
      fuse_type=FuseType.ABS_DIFF,
      **kwargs,
  ):
      super(RewardGenModel, self).__init__()
      if model_type == RewardGenModelType.GMNET:
          self.reward_gen_encoder = GMNet(
              in_dim, num_layers, hidden, pool_type=pool_type, **kwargs
          )
      elif model_type == RewardGenModelType.GENET:
          self.reward_gen_encoder = GENet(
              in_dim, num_layers, hidden, pool_type=pool_type, **kwargs
          )
  
      self.fuse_type = fuse_type
      if fuse_type == FuseType.CONCAT:
          self.pred_head = nn.Sequential(
              nn.Linear(2 * hidden, 2 * hidden),
              nn.ReLU(),
              nn.Linear(2 * hidden, 1),
              nn.Sigmoid(),
          )
      elif fuse_type == FuseType.COSINE:
          self.pred_head = nn.Sequential(
              nn.Linear(hidden, 2 * hidden),
              nn.ReLU(),
              nn.Linear(2 * hidden, hidden),
          )
          self.cos = torch.nn.CosineSimilarity(dim=1)
      else:
          in_hidden = hidden
          self.pred_head = nn.Sequential(
              nn.Linear(in_hidden, 2 * hidden),
              nn.ReLU(),
              nn.Linear(2 * hidden, 1),
              nn.Sigmoid(),
          )
  
file_path: DIG/dig/auggraph/method/GraphAug/model/reward_generator.py
imports: import torch
  import torch.nn
  from .genet import GENet
  from .gmnet import GMNet
  from dig.auggraph.method.GraphAug.constants import *
label_desc: Use image node_id 1 to create a new RewardGenModel object from inherited base classes: torch.nn with example: obj = RewardGenModel(in_dim, num_layers, hidden, pool_type, model_type, fuse_type)