code: def test_get_balanced_memory(self):
      model = ModelForTest()
      # model has size 236: linear1 64, batchnorm 72, linear2 100
      max_memory = get_balanced_memory(
          model, max_memory={0: 200, 1: 200}
      )
      self.assertDictEqual({0: 200, 1: 200}, max_memory)
  
      # We should be able to set models on a non-contiguous sub-set of
      max_memory = get_balanced_memory(
          model, max_memory={0: 200, 2: 200}
      )
      self.assertDictEqual({0: 200, 2: 200}, max_memory)
  
      max_memory = get_balanced_memory(
          model, max_memory={0: 300, 1: 300}
      )
      self.assertDictEqual({0: 215, 1: 300}, max_memory)
  
      # Last device always get max memory to give more buffer and avoid accidental CPU offload
      max_memory = get_balanced_memory(
          model, max_memory={0: 300, 1: 500}
      )
      self.assertDictEqual({0: 215, 1: 500}, max_memory)
  
      # Last device always get max memory to give more buffer, even if CPU is provided
      max_memory = get_balanced_memory(
          model, max_memory={0: 300, "cpu": 1000}
      )
      self.assertDictEqual({0: 300, "cpu": 1000}, max_memory)
  
      # If we set a device to 0, it's not counted.
      max_memory = get_balanced_memory(
          model, max_memory={0: 0, 1: 300, 2: 300}
      )
      self.assertDictEqual({0: 0, 1: 215, 2: 300}, max_memory)
  
      # If we set a device to 0, it's not counted.
      max_memory = get_balanced_memory(
          model, max_memory={0: 0, "cpu": 100}
      )
      self.assertDictEqual({0: 0, "cpu": 100}, max_memory)
  
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 23 for calling the ModelingUtilsTester obj's underlying member method code with example usage: obj.test_get_balanced_memory() without return types