Class FakeBothModel

tests/test_stateless_executor.py:220–232

    rewards = torch.tensor([0.2, 0.8])

    class FakeBothModel(torch.nn.Module):
        def __init__(self):
            super().__init__()
            self.calls = 0

        def forward(self, input_ids, attention_mask=None, use_cache=None):
            del attention_mask
            assert use_cache is False
            self.calls += 1
            return {
                "logits": logits[: input_ids.shape[0], : input_ids.shape[1]],
                "rewards": rewards,
            }

    model = FakeBothModel()
    executor = StatelessForwardExecutor(model, StatelessForwardConfig(mode="both"))