def compress(
        self,
        dataloader: _Optional[_Iterable] = None,
        loss_fn: _Optional[_Callable] = None,
        sensitivity_path: _Optional[str] = None,
        num_kmeans_workers: int = 1,
        num_sensitivity_workers: int = 1,
        inplace: bool = False,
        fsdp_auto_wrap_policy: _Optional[_FSDPAutoWrapPolicy] = None,
    ) -> _torch.nn.Module:
        """
        Compresses a model's weights using Fisher information sensitivity based weighted k-means
        palettization.

        Args:
            dataloader (:py:class:`Iterable`): An iterable where each element
                is an input to the model to be compressed. Used for computing gradients of model weights.
                This argument is not needed if ``sensitivity_path`` is specified and will be ignored.
                It is required then ``sensitivity_path`` is ``None``.  Defaults to ``None``.
            loss_fn (:obj:`Callable`): A callable which takes the model and data as input and performs
                a forward pass on the model and computes the training loss. This argument is not needed if
                ``sensitivity_path`` is specified and will be ignored. It is required when ``sensitivity_path``
                is ``None``. Defaults to ``None``.
            sensitivity_path (:obj:`str` or ``None``): An optional path from which the sensitivity values
                are loaded. If ``sensitivity_path`` is not ``None``, sensitivity values are loaded from the
                path specified, otherwise, sensitivity values are computed using the ``dataloader`` and
                ``loss_fn``. The sensitivity values stored at ``sensitivity_path`` should be a dictionary
                from strings indicating fully qualified parameter names to tensors with the same shape as the
                parameters, with each element of the tensor indicating how important that element is. This is
                usally the output of the :py:meth:`compute_sensitivity` method. Defaults to ``None``.
            num_kmeans_workers (:obj:`int`): Number of worker processes to use for performing k-means.
                It is recommended to use more than one worker process to parallelize the clustering,
                especially when multiple CPUs are available. Defaults to ``1``.
            num_sensitivity_workers (:obj:`int`): Number of worker processes to use for computing
                sensitivity. For large models, it is recommended to set this value to the number
                of GPUs available. Defaults to ``1``.
            inplace (:obj:`bool`): If ``True``, model transformations are carried out in-place and
                the original module is mutated, otherwise a copy of the model is mutated and returned.
                Defaults to ``False``.
            fsdp_auto_wrap_policy (:py:class:`_FSDPAutoWrapPolicy` or ``None``): Policy which specifies
                how different submodules of ``model`` are wrapped with individual
                :py:class:`FullyShardedDataParallel` wrappers. This argument is only used when
                ``num_sensitivity_workers > 1`` and it is only necessary when the model cannot be fit on a single GPU.
                Please refer to documentation of :py:class:`_FSDPAutoWrapPolicy` for more details.
                Defaults to ``None``.
        """
        self._model = super().compress(dataloader=dataloader, inplace=inplace)
        if sensitivity_path is None:
            sensitivity_dict = self.compute_sensitivity(
                dataloader,
                loss_fn,
                sensitivity_path,
                num_sensitivity_workers,
                fsdp_auto_wrap_policy=fsdp_auto_wrap_policy,
            )
        else:
            _logger.info(f"Loading sensitivity values from {sensitivity_path}.")
            sensitivity_dict = _torch.load(sensitivity_path)

        kmeans_config_dict = dict()
        for name, submodule in self._get_submodules_to_compress():
            submod_config = self._config.get_module_config(name, submodule)
            if submod_config is None:
                continue

            k_means_module_cls = _KMeansSupportedModulesRegistry.get_kmeans_module(submodule)
            if k_means_module_cls is None:
                continue

            for param_name in k_means_module_cls.parameter_names:
                # Validate configuration for parameter
                param = submodule.get_parameter(param_name)
                updated_config = _validate_param_config(
                    name + "." + param_name,
                    param,
                    submodule,
                    submod_config,
                    ["palettization_group_size", "palettization_cluster_dim"],
                )
                if not updated_config:
                    continue

                sensitivity_key = f"{name}.{param_name}" if len(name) > 0 else param_name
                sensitivity = sensitivity_dict[sensitivity_key]

                if name not in kmeans_config_dict:
                    kmeans_config_dict[name] = {}

                kmeans_config_dict[name][param_name] = _KMeansConfig(
                    n_bits=updated_config.n_bits,
                    axis=updated_config.channel_axis,
                    lut_dtype=updated_config.lut_dtype,
                    block_size=updated_config.group_size,
                    importance=sensitivity,
                    cluster_dim=updated_config.cluster_dim,
                    enable_per_channel_scale=updated_config.enable_per_channel_scale,
                )

        if num_kmeans_workers > 1:
            return _ParallelKMeans.cluster_weights(
                self._model, kmeans_config_dict, num_workers=num_kmeans_workers
            )
        else:
            return _SequentialKMeans.cluster_weights(self._model, kmeans_config_dict)

        Compresses a model's weights using Fisher information sensitivity based weighted k-means
        palettization.

        Args:
            dataloader (:py:class:`Iterable`): An iterable where each element
                is an input to the model to be compressed. Used for computing gradients of model weights.
                This argument is not needed if ``sensitivity_path`` is specified and will be ignored.
                It is required then ``sensitivity_path`` is ``None``.  Defaults to ``None``.
            loss_fn (:obj:`Callable`): A callable which takes the model and data as input and performs
                a forward pass on the model and computes the training loss. This argument is not needed if
                ``sensitivity_path`` is specified and will be ignored. It is required when ``sensitivity_path``
                is ``None``. Defaults to ``None``.
            sensitivity_path (:obj:`str` or ``None``): An optional path from which the sensitivity values
                are loaded. If ``sensitivity_path`` is not ``None``, sensitivity values are loaded from the
                path specified, otherwise, sensitivity values are computed using the ``dataloader`` and
                ``loss_fn``. The sensitivity values stored at ``sensitivity_path`` should be a dictionary
                from strings indicating fully qualified parameter names to tensors with the same shape as the
                parameters, with each element of the tensor indicating how important that element is. This is
                usally the output of the :py:meth:`compute_sensitivity` method. Defaults to ``None``.
            num_kmeans_workers (:obj:`int`): Number of worker processes to use for performing k-means.
                It is recommended to use more than one worker process to parallelize the clustering,
                especially when multiple CPUs are available. Defaults to ``1``.
            num_sensitivity_workers (:obj:`int`): Number of worker processes to use for computing
                sensitivity. For large models, it is recommended to set this value to the number
                of GPUs available. Defaults to ``1``.
            inplace (:obj:`bool`): If ``True``, model transformations are carried out in-place and
                the original module is mutated, otherwise a copy of the model is mutated and returned.
                Defaults to ``False``.
            fsdp_auto_wrap_policy (:py:class:`_FSDPAutoWrapPolicy` or ``None``): Policy which specifies
                how different submodules of ``model`` are wrapped with individual
                :py:class:`FullyShardedDataParallel` wrappers. This argument is only used when
                ``num_sensitivity_workers > 1`` and it is only necessary when the model cannot be fit on a single GPU.
                Please refer to documentation of :py:class:`_FSDPAutoWrapPolicy` for more details.
                Defaults to ``None``.
        
coremltools/optimize/torch/palettization/sensitive_k_means.py