docstr: 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.
docstr_tok: ["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","."]
code: def test_cpu_gpu_disk_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 we also add `disk` on the device_map.
      """
      device_map = {
          "transformer.word_embeddings": 0,
          "transformer.word_embeddings_layernorm": "cpu",
          "lm_head": 0,
          "transformer.h": 1,
          "transformer.ln_f": "disk",
      }
      bnb_config = BitsAndBytesConfig(
          llm_int8_enable_fp32_cpu_offload=True
      )
      with tempfile.TemporaryDirectory() as tmpdirname:
          # Load model
          model_8bit = AutoModelForCausalLM.from_pretrained(
              self.model_name,
              device_map=device_map,
              quantization_config=bnb_config,
              offload_folder=tmpdirname,
          )
  
          # 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", "disk"},
          )
  
          self.check_inference_correctness(model_8bit)
  
imports: import gc
  import importlib.metadata
  import tempfile
  import unittest
  from packaging import version
  from transformers import AutoConfig, AutoModel, AutoModelForCausalLM, AutoModelForSeq2SeqLM, AutoModelForSequenceClassification, AutoTokenizer, BitsAndBytesConfig, pipeline
  from transformers.testing_utils import is_accelerate_available, is_torch_available, require_accelerate, require_bitsandbytes, require_torch, require_torch_gpu, require_torch_multi_gpu, slow
label_desc: Use image node_id 5 for calling the MixedInt8TestCpuGpu obj's underlying member method code with example usage: obj.test_cpu_gpu_disk_loading_custom_device_map() without return types