def step(self):
        """
        Step through the palettizer. When the number of times ``step``
        is called is equal to ``milestone``, palettization is enabled.
        """
        for name, module in self._model.named_modules():
            # FSDP flattens model parameters leading to change in module names
            # _normalize_fsdp_module_name function call to get original module name
            normalized_name = _normalize_fsdp_module_name(name)
            if normalized_name in self._milestones:
                if self._step_count == self._milestones[normalized_name]:
                    self._enable_fake_palett_impl(module, True)
                    self._init_prune_threshold_and_module_wise_target_sparsity(module)
                if self._step_count > self._milestones[normalized_name]:
                    self._update_prune_threshold(module)
        self._step_count += 1

        Step through the palettizer. When the number of times ``step``
        is called is equal to ``milestone``, palettization is enabled.
        
coremltools/optimize/torch/palettization/palettizer.py