class_docstr: Represents a reward model, containing a reward head.
  Only a single reward is computed for each sequence.
docstr: :return: A dict with structure:
      reward: [batch, num_responses]
      eval_stats: dict of stats returned by eval_fn
code: def reward(
      self,
      query_tokens,
      response_tokens,
      eval_fn=None,
      eval_inputs=None,
      **kwargs,
  ):
      """
      :return: A dict with structure:
          reward: [batch, num_responses]
          eval_stats: dict of stats returned by eval_fn
      """
      last_response_indices = _response_indices(response_tokens).to(
          self.device
      )
      if self.task_hparams is not None:
          assert_eq(
              query_tokens.size(1), self.task_hparams.query.length
          )
          assert_eq(
              response_tokens.size(2), self.task_hparams.response.length
          )
      assert query_tokens.size(0) == response_tokens.size(0)
      if eval_fn is not None:
          eval_fn = _wrap_reward_model_fn(eval_fn)
          eval_inputs["last_response_index"] = last_response_indices
      result = self._eval(
          query_tokens,
          response_tokens,
          eval_fn=eval_fn,
          eval_inputs=eval_inputs,
          **kwargs,
      )
      result["reward"] = gather_one(
          result["reward"]["response"][:, :, 1:],
          last_response_indices,
          dim=2,
      )
  
      assert_shape_eq(
          result["reward"],
          (response_tokens.size(0), response_tokens.size(1)),
      )
      return result
  
file_path: summarize-from-feedback/summarize_from_feedback/reward_model.py
imports: import functools
  import torch
  from summarize_from_feedback import tasks
  from summarize_from_feedback.query_response_model import QueryResponseModel, PADDING_TOKEN
  from summarize_from_feedback.utils.torch_utils import first_true_indices, gather_one
  from summarize_from_feedback.utils.assertions import assert_shape_eq, assert_eq
label_desc: Use image node_id 2 for calling the RewardModel obj's underlying member method code with example usage: obj.reward(query_tokens, response_tokens, eval_fn, eval_inputs) and returns: result