def test_roberta_batching(self, device: str):
        """
        Checks that the batch of size 2 give twice the same results than the batch of size 1.
        """
        _, model = get_toy_model(device)
        sample = mk_sample("en", device, batch_size=1)
        slen = sample["net_input"]["src_lengths"][0]
        sample2 = mk_sample("en", device, batch_size=2)
        with torch.no_grad():
            z = model.encoder.forward(
                sample["net_input"]["src_tokens"], sample["net_input"]["src_lengths"]
            )
            z = z["encoder_out"][-1]
            logits, _ = model.forward(**sample["net_input"])

            z2 = model.encoder.forward(
                sample2["net_input"]["src_tokens"], sample["net_input"]["src_lengths"]
            )
            z2 = z2["encoder_out"][-1]
            logits2, _ = model.forward(**sample2["net_input"])

        self.assertEqual(z.shape, (slen, 1, 12))
        self.assertEqual(z2.shape, (slen, 2, 12))
        self.assertTensorEqual(logits2[0], logits2[1])
        self.assertTensorEqual(logits[0], logits2[0])

        Checks that the batch of size 2 give twice the same results than the batch of size 1.
        
edgelm/tests/test_roberta.py