code: def test_prepare_palettizer_simple_model_custom_palettization_config(
      simple_model,
  ):
      simple_model_copy = copy.deepcopy(simple_model)
      custom_config = {
          nn.Conv2d: {
              "n_bits": 2,
              "cluster_dim": 2,
              "kmeans_max_iter": 4,
          },
          nn.Linear: {
              "n_bits": 4,
              "cluster_dim": 1,
              "kmeans_max_iter": 5,
          },
      }
      config = DKMPalettizerConfig.from_dict(
          {"module_type_configs": custom_config}
      )
      palettizer = DKMPalettizer(simple_model, config)
      prepared_model = palettizer.prepare(inplace=True)
      num_epochs = 1
      for epoch in range(num_epochs):
          palettizer.step()
  
      _assert_changes_post_prepare(
          simple_model_copy.conv2,
          prepared_model.conv2,
          custom_config[nn.Conv2d]["n_bits"],
          custom_config[nn.Conv2d]["cluster_dim"],
          custom_config[nn.Conv2d]["kmeans_max_iter"],
      )
      _assert_changes_post_prepare(
          simple_model_copy.fc1,
          prepared_model.fc1,
          custom_config[nn.Linear]["n_bits"],
          custom_config[nn.Linear]["cluster_dim"],
          custom_config[nn.Linear]["kmeans_max_iter"],
      )
      _assert_changes_post_prepare(
          simple_model_copy.fc2,
          prepared_model.fc2,
          custom_config[nn.Linear]["n_bits"],
          custom_config[nn.Linear]["cluster_dim"],
          custom_config[nn.Linear]["kmeans_max_iter"],
      )
  
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 10 for calling a global function with example usage: test_prepare_palettizer_simple_model_custom_palettization_config(simple_model) without return types