def _scale_by_per_channel_scale(self, param_name: str, param: _torch.Tensor) -> _torch.Tensor:
        """
        Compute per channel scales for scaling the parameter in the range ``[-1, 1]``
        and store them in the parameter metadata. Also scale the parameter using
        the computed scales.
        """
        if self.config[param_name].enable_per_channel_scale:
            flattened_param = param.flatten(1)
            per_channel_scale = _torch.max(_torch.abs(flattened_param), dim=1, keepdim=True).values
            # Handle zero scales
            per_channel_scale[per_channel_scale == 0] = 1
            flattened_param /= per_channel_scale
            param = flattened_param.reshape(param.shape)
            self._parameter_metadata[param_name]["per_channel_scale"] = per_channel_scale
        return param

        Compute per channel scales for scaling the parameter in the range ``[-1, 1]``
        and store them in the parameter metadata. Also scale the parameter using
        the computed scales.
        
coremltools/optimize/torch/_utils/k_means.py