def _compute_sensitivity_impl_single_worker(
        self, dataset: _List, loss_fn: _Callable, sensitivity_path: _Optional[str]
    ):
        """
        Computes sensitivity for the model weights using a single process.
        """
        if _torch.cuda.is_available():
            self._model.cuda()

        self._model.zero_grad()

        with self._register_grad_square_hooks(self._model):
            for didx, data in enumerate(dataset):
                _logger.info(f"Computing sensitivity using sample {didx}")
                loss = loss_fn(self._model, data)
                loss.backward()

            sensitivity_dict = dict()
            for name, param in self._model.named_parameters(remove_duplicate=True):
                if param.requires_grad:
                    sensitivity_dict[name] = -param.grad.cpu()

            _torch.save(sensitivity_dict, self._get_sensitivity_path(sensitivity_path))

        Computes sensitivity for the model weights using a single process.
        
coremltools/optimize/torch/palettization/sensitive_k_means.py