def test_cpu_gpu_loading_custom_device_map(self):
        r"""
        A test to check is dispatching a model on cpu & gpu works correctly using a custom `device_map`.
        This time the device map is more organized than the test above and uses the abstraction
        `transformer.h` to encapsulate all the decoder layers.
        """
        device_map = {
            "transformer.word_embeddings": "cpu",
            "transformer.word_embeddings_layernorm": "cpu",
            "lm_head": "cpu",
            "transformer.h": 0,
            "transformer.ln_f": 1,
        }
        bnb_config = BitsAndBytesConfig(llm_int8_enable_fp32_cpu_offload=True)

        # Load model
        model_8bit = AutoModelForCausalLM.from_pretrained(
            self.model_name,
            device_map=device_map,
            quantization_config=bnb_config,
        )

        # Check that the model has been correctly set on device 0, 1, and `cpu`.
        self.assertEqual(set(model_8bit.hf_device_map.values()), {0, 1, "cpu"})

        self.check_inference_correctness(model_8bit)

        A test to check is dispatching a model on cpu & gpu works correctly using a custom `device_map`.
        This time the device map is more organized than the test above and uses the abstraction
        `transformer.h` to encapsulate all the decoder layers.
        
dpkd/transformers/tests/quantization/bnb/test_mixed_int8.py