code: def test_prepare_palettizer_different_milestone_per_module_type(
      simple_model,
  ):
      custom_config = {
          nn.Conv2d: {
              "n_bits": 2,
              "cluster_dim": 2,
              "kmeans_max_iter": 4,
              "milestone": 1,
          },
          nn.Linear: {
              "n_bits": 4,
              "kmeans_max_iter": 5,
              "milestone": 2,
          },
      }
      config = DKMPalettizerConfig.from_dict(
          {"module_type_configs": custom_config}
      )
      palettizer = DKMPalettizer(simple_model, config)
  
      orig_conv_mods = [simple_model.conv2]
      orig_fc_mods = [simple_model.fc1, simple_model.fc2]
  
      prepared_model = palettizer.prepare()
  
      prepared_conv_mods = [prepared_model.conv2]
      prepared_fc_mods = [prepared_model.fc1, prepared_model.fc2]
  
      num_steps = 3
      for step in range(num_steps):
          palettizer.step()
          if step == 0:
              for mod in prepared_conv_mods + prepared_fc_mods:
                  assert (
                      mod.weight_fake_quant.fake_palett_enabled[0] == 0
                  )
          elif step == 1:
              for mod in prepared_conv_mods:
                  assert (
                      mod.weight_fake_quant.fake_palett_enabled[0] == 1
                  )
              for mod in prepared_fc_mods:
                  assert (
                      mod.weight_fake_quant.fake_palett_enabled[0] == 0
                  )
          else:
              for mod in prepared_conv_mods:
                  assert (
                      mod.weight_fake_quant.fake_palett_enabled[0] == 1
                  )
              for mod in prepared_fc_mods:
                  assert (
                      mod.weight_fake_quant.fake_palett_enabled[0] == 1
                  )
  
      for orig, prep in zip(orig_conv_mods, prepared_conv_mods):
          _assert_changes_post_prepare(
              orig,
              prep,
              custom_config[nn.Conv2d]["n_bits"],
              custom_config[nn.Conv2d]["cluster_dim"],
              custom_config[nn.Conv2d]["kmeans_max_iter"],
          )
      for orig, prep in zip(orig_fc_mods, prepared_fc_mods):
          _assert_changes_post_prepare(
              orig,
              prep,
              custom_config[nn.Linear]["n_bits"],
              1,
              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 13 for calling a global function with example usage: test_prepare_palettizer_different_milestone_per_module_type(simple_model) without return types