def compute_sensitivity(
        self,
        dataloader: _Iterable,
        loss_fn: _Callable,
        sensitivity_path: _Optional[str] = None,
        num_sensitivity_workers: int = 1,
        fsdp_auto_wrap_policy: _Optional[_FSDPAutoWrapPolicy] = None,
    ) -> _Dict[str, _Any]:
        """
        Compute sensitivities of model's weights. A weight element's sensitivity indicates
        how much effect perturbing it has on the model's loss function. The sensitivities
        are computed as Fisher information of the model's weights.

        If ``sensitivity_path`` is passed as a non ``None`` value, the sensitivity matrices
        saved at the location specified by ``sensitivity_path``.

        When computing sensitivity of large models, it is recommended to use ``num_sensitivity_workers``
        equal to the number of GPUs available. This is because computing gradients using a single
        process maybe infeasible for a large model as it may not fit on a single GPU.
        When ``num_sensitivity_workers > 1``, the model is sharded on multiple GPUs using
        :py:class:`FullyShardedDataParallel`, which enables distributed computation of gradients.

        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.
            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
            sensitivity_path (:obj:`str` or ``None``): An optional path for saving the sensitivity
                of weights. Defaults to ``None``.
            num_sensitivity_workers (:obj:`int`): Number of worker processes used for computing sensitivity.
                Defaults to ``1``.
            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``.
        """
        if num_sensitivity_workers > 1 and not _torch.cuda.is_available():
            _logger.warning(
                "num_sensitivity_workers > 1 is only supported on GPUs with CUDA. Setting "
                "num_sensitivity_workers to 1, since a CUDA compatible PyTorch installation"
                "couldn't be found."
            )
            num_sensitivity_workers = 1

        # We save the model's state dict so that we can restore it later
        # We need to do this because _compute_sensitivity_impl_multiple_workers
        # sets the parameters' value to squares of their gradients and
        # _compute_sensitivity_impl_single_worker can modify layers such as batch norm
        # during forward pass
        _torch.save(self._model.state_dict(), self._model_checkpoint_path)
        if num_sensitivity_workers == 1:
            self._compute_sensitivity_impl_single_worker(
                self._get_dataset(0, 1, dataloader),
                loss_fn,
                sensitivity_path,
            )
        else:
            if fsdp_auto_wrap_policy is None:
                _logger.warning(
                    "num_sensitivity_workers > 1 and fsdp_auto_wrap_policy is None. For a large model, this might "
                    "lead to OOM issue on GPUs. Consider setting fsdp_auto_wrap_policy to indicate how different "
                    "submodules of the model should be wrapped with FSDP wrappers to prevent all gather for all "
                    "parameters on all GPUs."
                )

            ctx = _mp.get_context("spawn")

            worker_processes = [
                ctx.Process(
                    target=self._compute_sensitivity_impl_multiple_workers,
                    args=(
                        rank,
                        num_sensitivity_workers,
                        self._get_dataset(rank, num_sensitivity_workers, dataloader),
                        loss_fn,
                        sensitivity_path,
                        fsdp_auto_wrap_policy,
                    ),
                    name=f"Process-{rank}",
                )
                for rank in range(num_sensitivity_workers)
            ]
            for worker_process in worker_processes:
                worker_process.start()
                _logger.info(f"Started {worker_process.name} for computing sensitivity.")

            for worker_process in worker_processes:
                worker_process.join()
                _logger.info(f"Finished {worker_process.name}.")

        # restore the original state of the model
        self._model.cpu()
        old_state_dict = _torch.load(self._model_checkpoint_path)
        self._model.load_state_dict(old_state_dict)

        return self._process_sensitivity(sensitivity_path)

        Compute sensitivities of model's weights. A weight element's sensitivity indicates
        how much effect perturbing it has on the model's loss function. The sensitivities
        are computed as Fisher information of the model's weights.

        If ``sensitivity_path`` is passed as a non ``None`` value, the sensitivity matrices
        saved at the location specified by ``sensitivity_path``.

        When computing sensitivity of large models, it is recommended to use ``num_sensitivity_workers``
        equal to the number of GPUs available. This is because computing gradients using a single
        process maybe infeasible for a large model as it may not fit on a single GPU.
        When ``num_sensitivity_workers > 1``, the model is sharded on multiple GPUs using
        :py:class:`FullyShardedDataParallel`, which enables distributed computation of gradients.

        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.
            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
            sensitivity_path (:obj:`str` or ``None``): An optional path for saving the sensitivity
                of weights. Defaults to ``None``.
            num_sensitivity_workers (:obj:`int`): Number of worker processes used for computing sensitivity.
                Defaults to ``1``.
            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