Method forward_logits

rl_engine/alignment/model_wrappers.py:87–88

(self, input_ids: torch.Tensor, **model_kwargs: Any)

        return self.model(input_ids, **model_kwargs)

    def forward_logits(self, input_ids: torch.Tensor, **model_kwargs: Any) -> torch.Tensor:
        return extract_logits(self.forward(input_ids, **model_kwargs))

    def selected_logprobs(
        self,