code: def forward(self, data1, data2):
      embed1, embed2 = self.reward_gen_encoder(data1, data2)
  
      if self.fuse_type == FuseType.ADD:
          pair_embed = embed1 + embed2
      elif self.fuse_type == FuseType.MULTIPLY:
          pair_embed = embed1 * embed2
      elif self.fuse_type == FuseType.CONCAT:
          pair_embed = torch.cat((embed1, embed2), dim=1)
      elif self.fuse_type == FuseType.ABS_DIFF:
          pair_embed = torch.abs(embed1 - embed2)
      elif self.fuse_type == FuseType.COSINE:
          embed1, embed2 = self.pred_head(embed1), self.pred_head(
              embed2
          )
          prob = (1.0 + self.cos(embed1, embed2)) / 2.0
          return prob
  
      prob = self.pred_head(pair_embed)
      return prob
  
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 2 for calling the RewardGenModel obj's underlying member method code with example usage: obj.forward(data1, data2) and returns: prob, prob