Method __init__

benchmarks/benchmark_stateless_executor.py:62–66

(self, vocab_size: int, hidden_dim: int)

    """Small deterministic causal-LM-like module for executor smoke benchmarks."""

    def __init__(self, vocab_size: int, hidden_dim: int):
        super().__init__()
        self.embedding = torch.nn.Embedding(vocab_size, hidden_dim)
        self.proj = torch.nn.Linear(hidden_dim, vocab_size)
        self.use_cache_calls: list[bool | None] = []

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