Method __init__

tests/test_paged_kv_baseline.py:22–25

(self, logits: torch.Tensor)

class FakeGenerationReferenceModel(torch.nn.Module):
    def __init__(self, logits: torch.Tensor):
        super().__init__()
        self.register_buffer("fixed_logits", logits)
        self.use_cache_calls: list[bool | None] = []

    def forward(self, input_ids, attention_mask=None, use_cache=None):
        del attention_mask