code: def compute_rewards(
      self, prompts, log_probs, ref_log_probs, reward_score, action_mask
  ):
      kl_divergence_estimate = -self.kl_ctl * (
          log_probs - ref_log_probs
      )
      rewards = kl_divergence_estimate
      start = prompts.shape[1] - 1
      ends = start + action_mask[:, start:].sum(1) + 1
      reward_clip = torch.clamp(
          reward_score, -self.clip_reward_value, self.clip_reward_value
      )
      batch_size = log_probs.shape[0]
      for j in range(batch_size):
          rewards[j, start : ends[j]][-1] += reward_clip[j]
  
      return rewards
  
imports: import torch
  import torch.nn.functional
  import time
  import deepspeed
  from deepspeed.runtime.zero.partition_parameters import ZeroParamStatus
  from deepspeed.accelerator import get_accelerator
  from dschat.utils.utils import print_rank_0
label_desc: Use image node_id 4 for calling the DeepSpeedPPOTrainer obj's underlying member method code with example usage: obj.compute_rewards(prompts, log_probs, ref_log_probs, reward_score, action_mask) and returns: rewards