code: def _get_predicted_reward(self, batch: rlt.BanditRewardModelInput):
      model_rewards_all_actions = self.reward_net(batch.state)
      logged_action_idxs = torch.argmax(
          batch.action, dim=1, keepdim=True
      )
      predicted_reward = model_rewards_all_actions.gather(
          1, logged_action_idxs
      )
      return predicted_reward
  
imports: import logging
  from typing import Optional
  import numpy
  import reagent.core.types
  import torch
  from reagent.core.dataclasses import field
  from reagent.models.base import ModelBase
  from reagent.optimizer.union import Optimizer__Union
  from reagent.training.reagent_lightning_module import ReAgentLightningModule
  from reagent.training.reward_network_trainer import _get_loss_function, LossFunction
label_desc: Use image node_id 4 for calling the BanditRewardNetTrainer obj's underlying member method code with example usage: obj._get_predicted_reward(batch) and returns: predicted_reward