code: def test_cpu_gpu_disk_loading_custom_device_map_kwargs(self):
      from transformers import AutoConfig, AutoModelForCausalLM
  
      r"""
          A test to check is dispatching a model on cpu & gpu works correctly using a custom `device_map`.
          This time we also add `disk` on the device_map - using the kwargs directly instead of the quantization config
          """
      device_map = {
          "transformer.word_embeddings": 0,
          "transformer.word_embeddings_layernorm": "disk",
          "lm_head": 0,
          "transformer.h": 1,
          "transformer.ln_f": "cpu",
      }
      bnb_quantization_config = BnbQuantizationConfig(load_in_4bit=True)
  
      with init_empty_weights():
          model_4bit = AutoModelForCausalLM.from_config(
              AutoConfig.from_pretrained(self.model_name)
          )
  
      model_4bit.tie_weights()
      with tempfile.TemporaryDirectory() as tmpdirname:
          model_4bit = load_and_quantize_model(
              model_4bit,
              bnb_quantization_config,
              weights_location=self.weights_location,
              device_map=device_map,
              no_split_module_classes=["BloomBlock"],
              offload_folder=tmpdirname,
              offload_state_dict=True,
          )
          self.check_inference_correctness(model_4bit)
  
imports: import gc
  import tempfile
  import unittest
  import torch
  import torch.nn
  from accelerate import Accelerator, init_empty_weights
  from accelerate.test_utils import require_bnb, require_cuda, require_huggingface_suite, require_multi_gpu, slow
  from accelerate.utils.bnb import load_and_quantize_model
  from accelerate.utils.dataclasses import BnbQuantizationConfig
label_desc: Use image node_id 10 for calling the Bnb4BitEmptyModelTest obj's underlying member method code with example usage: obj.test_cpu_gpu_disk_loading_custom_device_map_kwargs() without return types