def test_small_logits(self):
        r"""
        Logits testing to check implementation consistency between `t5x` implementation
        and `transformers` implementation of Switch-C transformers. We only check the logits
        of the first batch.
        """
        model = SwitchTransformersModel.from_pretrained("google/switch-base-8", torch_dtype=torch.bfloat16).to(
            torch_device
        )
        input_ids = torch.ones((32, 64), dtype=torch.long).to(torch_device)
        decoder_input_ids = torch.ones((32, 64), dtype=torch.long).to(torch_device)

        # fmt: off
        EXPECTED_MEAN_LOGITS = torch.Tensor(
            [
                -0.204102, -0.193359, 0.523438, -0.296875, 0.108887,
                0.0211182, 0.605469, -0.100586, -0.0551758, 0.296875,
                0.0090332, 0.174805, 0.139648, -0.170898, -0.0981445,
                0.0245361, 0.0373535, 0.050293, -0.212891, 0.129883,
                0.390625, -0.203125, -0.122559, -0.180664, 0.0437012,
                -0.349609, -0.0250244, -0.104004, -0.15918, -0.133789
            ]
        ).to(torch.bfloat16)
        # fmt: on
        hf_logits = model(input_ids, decoder_input_ids=decoder_input_ids).last_hidden_state.cpu()
        hf_logits = hf_logits[0, 0, :30]

        torch.testing.assert_allclose(hf_logits, EXPECTED_MEAN_LOGITS, rtol=6e-3, atol=9e-3)

        Logits testing to check implementation consistency between `t5x` implementation
        and `transformers` implementation of Switch-C transformers. We only check the logits
        of the first batch.
        
dpkd/transformers/tests/models/switch_transformers/test_modeling_switch_transformers.py