code: def test_infer_auto_device_map(self):
      model = ModelForTest()
      # model has size 236: linear1 64, batchnorm 72, linear2 100
  
      device_map = infer_auto_device_map(
          model, max_memory={0: 200, 1: 200}
      )
      # only linear1 fits on device 0 as we keep memory available for the maximum layer in case of offload
      self.assertDictEqual(
          device_map, {"linear1": 0, "batchnorm": 1, "linear2": 1}
      )
  
      device_map = infer_auto_device_map(
          model, max_memory={0: 200, 1: 172, 2: 200}
      )
      # On device 1, we don't care about keeping size available for the max layer, so even if there is just the
      # size available for batchnorm + linear2, they fit here.
      self.assertDictEqual(
          device_map, {"linear1": 0, "batchnorm": 1, "linear2": 1}
      )
  
      model.linear1.weight = model.linear2.weight
      device_map = infer_auto_device_map(
          model, max_memory={0: 200, 1: 200}
      )
      # By tying weights, the whole model fits on device 0
      self.assertDictEqual(device_map, {"": 0})
  
      # When splitting a bigger model, the split is done at the layer level
      model = nn.Sequential(
          ModelForTest(), ModelForTest(), ModelForTest()
      )
      device_map = infer_auto_device_map(
          model, max_memory={0: 500, 1: 500}
      )
      self.assertDictEqual(
          device_map,
          {
              "0": 0,
              "1.linear1": 0,
              "1.batchnorm": 0,
              "1.linear2": 1,
              "2": 1,
          },
      )
  
      # With no_split_module_classes, it's done at that module level
      model = nn.Sequential(
          ModelForTest(), ModelForTest(), ModelForTest()
      )
      device_map = infer_auto_device_map(
          model,
          max_memory={0: 500, 1: 500},
          no_split_module_classes=["ModelForTest"],
      )
      self.assertDictEqual(device_map, {"0": 0, "1": 1, "2": 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 20 for calling the ModelingUtilsTester obj's underlying member method code with example usage: obj.test_infer_auto_device_map() without return types