code: def test_load_checkpoint_in_model_two_gpu(self):
      device_map = {"linear1": 0, "batchnorm": "cpu", "linear2": 1}
  
      # Check with whole checkpoint
      model = ModelForTest()
      with tempfile.TemporaryDirectory() as tmp_dir:
          fname = os.path.join(tmp_dir, "pt_model.bin")
          torch.save(model.state_dict(), fname)
          load_checkpoint_in_model(model, fname, device_map=device_map)
      self.assertEqual(model.linear1.weight.device, torch.device(0))
      self.assertEqual(
          model.batchnorm.weight.device, torch.device("cpu")
      )
      self.assertEqual(model.linear2.weight.device, torch.device(1))
  
      # Check with sharded index
      model = ModelForTest()
      with tempfile.TemporaryDirectory() as tmp_dir:
          self.shard_test_model(model, tmp_dir)
          index_file = os.path.join(tmp_dir, "weight_map.index.json")
          load_checkpoint_in_model(
              model, index_file, device_map=device_map
          )
  
      self.assertEqual(model.linear1.weight.device, torch.device(0))
      self.assertEqual(
          model.batchnorm.weight.device, torch.device("cpu")
      )
      self.assertEqual(model.linear2.weight.device, torch.device(1))
  
      # Check with sharded checkpoint
      model = ModelForTest()
      with tempfile.TemporaryDirectory() as tmp_dir:
          self.shard_test_model(model, tmp_dir)
          load_checkpoint_in_model(
              model, tmp_dir, device_map=device_map
          )
  
      self.assertEqual(model.linear1.weight.device, torch.device(0))
      self.assertEqual(
          model.batchnorm.weight.device, torch.device("cpu")
      )
      self.assertEqual(model.linear2.weight.device, torch.device(1))
  
imports: import json
  import os
  import tempfile
  import unittest
  from collections import OrderedDict
  import torch
  import torch.nn
  from safetensors.torch import save_file
  from accelerate import init_empty_weights
  from accelerate.test_utils import require_cuda, require_huggingface_suite, require_multi_gpu
  from accelerate.utils.modeling import check_device_map, clean_device_map, compute_module_sizes, convert_file_size_to_int, find_tied_parameters, get_balanced_memory, infer_auto_device_map, load_checkpoint_in_model, load_state_dict, named_module_tensors, retie_parameters, set_module_tensor_to_device
label_desc: Use image node_id 17 for calling the ModelingUtilsTester obj's underlying member method code with example usage: obj.test_load_checkpoint_in_model_two_gpu() without return types