def _register_activation_processing_hook(
        self, atomic_layer, compressor_obj
    ) -> _torch.utils.hooks.RemovableHandle:
        """
        Registers a forward hook on the layer for performing computation
        using the inputs to acquire statistics. Returns the handle for
        the forward hook
        """

        def activation_processing_hook(_, inp, out):
            compressor_obj.add_batch(inp[0].data, out.data)

        return atomic_layer.register_forward_hook(activation_processing_hook)

        Registers a forward hook on the layer for performing computation
        using the inputs to acquire statistics. Returns the handle for
        the forward hook
        
coremltools/optimize/torch/layerwise_compression/layerwise_compressor.py