(self, vocab_size: int, hidden_dim: int)
| 79 | """Small deterministic scalar-reward module for executor smoke benchmarks.""" |
| 80 | |
| 81 | def __init__(self, vocab_size: int, hidden_dim: int): |
| 82 | super().__init__() |
| 83 | self.embedding = torch.nn.Embedding(vocab_size, hidden_dim) |
| 84 | self.reward_head = torch.nn.Linear(hidden_dim, 1) |
| 85 | self.use_cache_calls: list[bool | None] = [] |
| 86 | |
| 87 | def forward(self, input_ids, attention_mask=None, use_cache=None): |
| 88 | self.use_cache_calls.append(use_cache) |