def test_bnb_4bit_errors_loading_incorrect_state_dict(self):
        r"""
        Test if loading with an incorrect state dict raises an error.
        """
        with tempfile.TemporaryDirectory() as tmpdirname:
            nf4_config = BitsAndBytesConfig(load_in_4bit=True)
            model_4bit = SD3Transformer2DModel.from_pretrained(
                self.model_name, subfolder="transformer", quantization_config=nf4_config, device_map=torch_device
            )
            model_4bit.save_pretrained(tmpdirname)
            del model_4bit

            with self.assertRaises(ValueError) as err_context:
                state_dict = safetensors.torch.load_file(
                    os.path.join(tmpdirname, "diffusion_pytorch_model.safetensors")
                )

                # corrupt the state dict
                key_to_target = "context_embedder.weight"  # can be other keys too.
                compatible_param = state_dict[key_to_target]
                corrupted_param = torch.randn(compatible_param.shape[0] - 1, 1)
                state_dict[key_to_target] = bnb.nn.Params4bit(corrupted_param, requires_grad=False)
                safetensors.torch.save_file(
                    state_dict, os.path.join(tmpdirname, "diffusion_pytorch_model.safetensors")
                )

                _ = SD3Transformer2DModel.from_pretrained(tmpdirname)

            assert key_to_target in str(err_context.exception)

        Test if loading with an incorrect state dict raises an error.
        
tests/quantization/bnb/test_4bit.py