def cache(
        self, dataloader: _Iterable, nsamples: int, device: str
    ) -> _Tuple[_List[_torch.Tensor], _Dict[str, _torch.Tensor]]:
        """
        Cache inputs and keyword arguments to be fed to the first decoder layer
        of a GPT style model.

        Args:
            dataloader (:py:class:`Iterable`): An iterable where each element
                is an input to the model to be compressed.
            nsamples (:obj:`int`): Number of samples to cache.
            device (:obj:`str`): Device string for device to run compression on.
        """
        for layer in self._pre_layers:
            layer.to(device)

        inputs, kwarg_inputs = [], {}
        input_cacher_handle = self._first_layer.register_forward_pre_hook(
            self._get_input_cacher_pre_hook(inputs, kwarg_inputs), with_kwargs=True
        )
        self._feed_data(dataloader, nsamples, device)
        input_cacher_handle.remove()

        for layer in self._pre_layers:
            layer.cpu()

        for key, val in kwarg_inputs.items():
            if isinstance(val, _torch.Tensor):
                kwarg_inputs[key] = val.to(device)

        return inputs, kwarg_inputs

        Cache inputs and keyword arguments to be fed to the first decoder layer
        of a GPT style model.

        Args:
            dataloader (:py:class:`Iterable`): An iterable where each element
                is an input to the model to be compressed.
            nsamples (:obj:`int`): Number of samples to cache.
            device (:obj:`str`): Device string for device to run compression on.
        
coremltools/optimize/torch/layerwise_compression/input_cacher.py