def _get_input_cacher_pre_hook(inputs, kwarg_inputs):
        """
        Returns forward_pre_hook for caching inputs and keyword arguments
        to the first decoder layer of a GPT model.
        """

        def input_cacher_pre_hook(module, args, kwargs):
            inputs.append(args)
            for key, val in kwargs.items():
                kwarg_inputs[key] = val
            raise StopExecution()

        return input_cacher_pre_hook

        Returns forward_pre_hook for caching inputs and keyword arguments
        to the first decoder layer of a GPT model.
        
coremltools/optimize/torch/layerwise_compression/input_cacher.py