def set_param_shape(self, shape: _List[int]):
        """
        Method that sets the parameter shape. This method should only be called before starting training
        and after calling the set_ch_axis method.
        Should only be used for per-channel quantization.

        Args:
            shape (:obj:`List[int]`): Extent of the learnable quantization parameters.
        """

        assert self.ch_axis != -1

        self.shape = shape
        self.mask_shape = [1] * len(self.shape)
        self.mask_shape[self.ch_axis] = self.shape[self.ch_axis]
        self.channel_count = self.mask_shape[self.ch_axis]

        if self.device is None:
            self.scale = _torch.nn.Parameter(_torch.ones(self.mask_shape))
            self.zero_point = _torch.nn.Parameter(
                _torch.zeros(self.mask_shape),
                requires_grad=not _is_symmetric_quant(self.qscheme),
            )
        else:
            self.scale = _torch.nn.Parameter(_torch.ones(self.mask_shape, device=self.device))
            self.zero_point = _torch.nn.Parameter(
                _torch.zeros(self.mask_shape, device=self.device),
                requires_grad=not _is_symmetric_quant(self.qscheme),
            )

        Method that sets the parameter shape. This method should only be called before starting training
        and after calling the set_ch_axis method.
        Should only be used for per-channel quantization.

        Args:
            shape (:obj:`List[int]`): Extent of the learnable quantization parameters.
        
coremltools/optimize/torch/quantization/modules/learnable_fake_quantize.py