    def test_ngram_blocking_same_as_default_implem(self):
        """Test that cuda extension returns same things as default impl in many settings."""
        vocab_size = 4
        step = 6
        for _ in range(2):
            block_param = np.random.choice([1, 2, 3, 4])
            batch_size = np.random.randint(1, 8)
            beam_size = np.random.choice([1, 2, 4, 8])
            lprobs = torch.zeros((beam_size * batch_size, vocab_size), device="cuda")

            generated_tok = torch.tensor(
                np.random.randint(
                    0, vocab_size, size=(batch_size * beam_size, step + 1)
                ),
                device="cuda",
                dtype=torch.long,
            )
            self._compare_cuda_ext_to_default_implem(
                batch_size,
                beam_size,
                generated_tok,
                lprobs,
                step,
                block_param,
            )
Test that cuda extension returns same things as default impl in many settings.
tests/test_sequence_generator.py