code: def test_attach_config_simple_model_custom_palettization_config(
      simple_model,
  ):
      custom_config = {
          nn.Conv2d: {"n_bits": 2, "cluster_dim": 2},
          nn.Linear: {"n_bits": 4},
      }
      config = DKMPalettizerConfig.from_dict(
          {
              "module_type_configs": custom_config,
              "module_name_configs": {
                  "conv2": {"n_bits": 3, "cluster_dim": 2}
              },
          }
      )
      palettizer = DKMPalettizer(simple_model, config)
      palettizer.prepare(inplace=True)
  
      _assert_changes_post_attach(simple_model.conv2, 3, 2)
      _assert_changes_post_attach(
          simple_model.fc1, custom_config[nn.Linear]["n_bits"], 1
      )
      _assert_changes_post_attach(
          simple_model.fc2, custom_config[nn.Linear]["n_bits"], 1
      )
  
imports: import copy
  import pytest
  import torch
  import torch.functional
  import torch.nn
  from coremltools.optimize.torch.palettization import DKMPalettizer, DKMPalettizerConfig
  from coremltools.optimize.torch.palettization.palettization_config import DEFAULT_PALETTIZATION_SCHEME
  from coremltools.test.optimize.torch.palettization.palettization_utils import _assert_changes_post_attach, _assert_changes_post_prepare
label_desc: Use image node_id 5 for calling a global function with example usage: test_attach_config_simple_model_custom_palettization_config(simple_model) without return types