code: def _wrap_reward_model_fn(fn):
      @functools.wraps(fn)
      def wrapped(outputs_mb, inputs_mb):
          rewards = outputs_mb["reward"]["response"][:, :, 1:]
          rewards = gather_one(
              rewards, inputs_mb["last_response_index"], dim=2
          )
          outputs_mb["reward"] = rewards
          return fn(outputs_mb, inputs_mb)
  
      return wrapped
  
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 a global function with example usage: _wrap_reward_model_fn(fn) and returns: wrapped, fn