def test_roberta_incremental_decoder(self, device: str):
        """
        Checks that incremental decoding yields the same result than non incremental one.
        """
        task, model = get_toy_model(device)

        en_sample = mk_sample("en", device)
        en_tokens = en_sample["net_input"]["src_tokens"]
        ro_sample = mk_sample("ro", device)
        ro_tokens = ro_sample["net_input"]["src_tokens"]

        en_enc = model.encoder.forward(
            en_tokens, src_lengths=en_sample["net_input"]["src_lengths"]
        )
        (bs, tgt_len) = ro_tokens.shape

        # Decode without incremental state
        ro_dec, _ = model.decoder.forward(ro_tokens, encoder_out=en_enc)
        self.assertEqual(ro_dec.shape, (bs, tgt_len, VOCAB_SIZE))
        self.assertTensorEqual(ro_dec[0], ro_dec[1])

        # Decode with incremental state
        inc_state = {}
        ro_dec_inc = []
        for l in range(tgt_len):
            ro, _ = model.decoder.forward(
                ro_tokens[:, : l + 1], encoder_out=en_enc, incremental_state=inc_state
            )
            self.assertEqual(ro.shape, (bs, 1, VOCAB_SIZE))
            ro_dec_inc.append(ro)

        for l in range(tgt_len):
            # Intra-batch
            self.assertTensorEqual(ro_dec_inc[l][0], ro_dec_inc[l][1])
            # Incremental vs non-incremental
            self.assertTensorEqual(ro_dec_inc[l][:, 0], ro_dec[:, l])

        Checks that incremental decoding yields the same result than non incremental one.
        
edgelm/tests/test_roberta.py