def _construct_compression_info(
        state_dict: Dict[str, torch.Tensor],
    ) -> Dict[str, CompressionInfo]:
        """
        Construct compression info from the traced model's state_dict.

        The state_dict of the traced model is something like
        {
            'dense1.weight': xxx, 'dense1.bias': xxx,
            'dense1._COREML_/weight/quantization_n_bits': tensor(4),
            'dense1._COREML_/weight/quantization_scale': xxx,
            'dense1._COREML_/weight/zero_point': xxx,
            'dense1._COREML_/weight/compression_type': tensor([3]),
            'dense2.weight': xxx,
            ...
        }

        We extract the compression info and store it as a dict
        {
            'dense1.weight': CompressionInfo(quantization_n_bits=4, quantization_scale=xxx,
                                            zero_point=xxx, compression_type=[QUANTIZATION]),
            'dense2.weight': ...
        }
        """
        compression_info = dict()
        for torch_key_name in state_dict.keys():
            if f"{_COMPRESSION_INFO_PREFIX}/metadata_version" in torch_key_name:
                # TODO: rdar://124707382 ([Compression] Support versioning in CompressionInfo)
                continue

            if _COMPRESSION_INFO_PREFIX in torch_key_name:
                module_name = None
                buffer_name = torch_key_name
                if not torch_key_name.startswith(_COMPRESSION_INFO_PREFIX):
                    module_name, buffer_name = torch_key_name.rsplit(".", 1)
                _, param_name, compression_key = buffer_name.rsplit("/", 2)
                if module_name:
                    param_name = f"{module_name}.{param_name}"

                if param_name not in compression_info:
                    compression_info[param_name] = CompressionInfo()
                setattr(
                    compression_info[param_name],
                    compression_key,
                    state_dict[torch_key_name],
                )

        return compression_info

        Construct compression info from the traced model's state_dict.

        The state_dict of the traced model is something like
        {
            'dense1.weight': xxx, 'dense1.bias': xxx,
            'dense1._COREML_/weight/quantization_n_bits': tensor(4),
            'dense1._COREML_/weight/quantization_scale': xxx,
            'dense1._COREML_/weight/zero_point': xxx,
            'dense1._COREML_/weight/compression_type': tensor([3]),
            'dense2.weight': xxx,
            ...
        }

        We extract the compression info and store it as a dict
        {
            'dense1.weight': CompressionInfo(quantization_n_bits=4, quantization_scale=xxx,
                                            zero_point=xxx, compression_type=[QUANTIZATION]),
            'dense2.weight': ...
        }
        
coremltools/converters/mil/frontend/torch/converter.py