Method __init__

tests/test_alignment_model_wrappers.py:21–25

(self, logits: torch.Tensor, *, output_kind: str = "tensor")

class FixedLogitsModel(torch.nn.Module):
    def __init__(self, logits: torch.Tensor, *, output_kind: str = "tensor"):
        super().__init__()
        self.output_kind = output_kind
        self.weight = torch.nn.Parameter(torch.tensor(1.0))
        self.register_buffer("base_logits", logits.clone())

    def forward(self, input_ids: torch.Tensor, *, logit_bias: float = 0.0):
        del input_ids