def _compare_cuda_ext_to_default_implem(
        self, bsz, beam_size, generated_tok, lprobs, step, block_param
    ):
        """Assert that cuda extension and default implem return the same thing."""
        blocker = NGramRepeatBlock(block_param)
        assert blocker.use_extension, "Extension not compiled"
        cuda_ext_result = blocker(
            generated_tok,
            lprobs.clone(),
            bsz,
            beam_size,
            step,
        )
        blocker.use_extension = False
        baseline_result = blocker(
            generated_tok,
            lprobs.clone(),
            bsz,
            beam_size,
            step,
        )
        self.assertTensorEqual(cuda_ext_result, baseline_result)
        blocker.use_extension = True
        return cuda_ext_result, baseline_result
Assert that cuda extension and default implem return the same thing.
edgelm/tests/test_sequence_generator.py