class_docstr: Reward model base class.
  
  Args:
      model (nn.Module): Reward model.
      value_head (nn.Module): Value head to get reward score.
      lora_rank (int): LoRA rank.
      lora_train_bias (str): LoRA bias training mode.
code: def forward(
      self, sequences: torch.LongTensor, attention_mask: torch.Tensor
  ) -> torch.Tensor:
      outputs = self.model(sequences, attention_mask=attention_mask)
      last_hidden_states = outputs["last_hidden_state"]
      sequence_lengths = torch.max(
          attention_mask
          * torch.arange(sequences.size(1), device=sequences.device),
          dim=1,
      )[0]
      sequence_hidden_states = last_hidden_states[
          torch.arange(last_hidden_states.size(0)), sequence_lengths
      ]
      values = self.value_head(sequence_hidden_states).squeeze(
          1
      )  # ensure shape is (B, )
      return values
  
file_path: colossalai/applications/Chat/coati/models/base/reward_model.py
imports: from typing import Optional
  import torch
  import torch.nn
  from ..lora import LoRAModule
label_desc: Use image node_id 2 for calling the RewardModel obj's underlying member method code with example usage: obj.forward(sequences, attention_mask) and returns: values