Method compact_token_ids

rl_engine/testing/rl_batch.py:67–68

(self)

        return dense.reshape(-1, *dense.shape[2:])[self.flat_completion_mask]

    def compact_token_ids(self) -> torch.Tensor:
        return self.flat_token_ids[self.flat_completion_mask]

    def benchmark_metadata(self) -> dict[str, Any]:
        return {