code: def test_attach_config_weight_threshold_range_different_milestone(
      simple_model,
  ):
      custom_config = {
          nn.Conv2d: [
              {
                  "n_bits": 4,
                  "cluster_dim": 4,
                  "weight_threshold": 1000,
                  "milestone": 2,
              },
              {
                  "n_bits": 2,
                  "cluster_dim": 2,
                  "weight_threshold": 400,
                  "milestone": 1,
              },
          ]
      }
      config = DKMPalettizerConfig.from_dict(
          {"module_type_configs": custom_config}
      )
      palettizer = DKMPalettizer(simple_model, config)
      prepared_model = palettizer.prepare()
  
      # configs should get sorted automatically
      assert (
          hasattr(prepared_model.fc1, "qconfig")
          and prepared_model.fc1.qconfig is None
      )
  
      num_steps = 3
      for step in range(num_steps):
          palettizer.step()
          if step == 0:
              assert (
                  prepared_model.conv2.weight_fake_quant.fake_palett_enabled[
                      0
                  ]
                  == 0
              )
          elif step == 1:
              assert (
                  prepared_model.conv2.weight_fake_quant.fake_palett_enabled[
                      0
                  ]
                  == 0
              )
          else:
              assert (
                  prepared_model.conv2.weight_fake_quant.fake_palett_enabled[
                      0
                  ]
                  == 1
              )
  
      _assert_changes_post_attach(
          prepared_model.conv2,
          custom_config[nn.Conv2d][0]["n_bits"],
          custom_config[nn.Conv2d][0]["cluster_dim"],
      )
  
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 14 for calling a global function with example usage: test_attach_config_weight_threshold_range_different_milestone(simple_model) without return types