def set_ch_axis(self, ch_axis: int):
        """
        Method that sets the channel axis. This method should only be called before starting training.

        Args:
            ch_axis (:obj:`int`): Non-negative channel axis index.
        """

        assert ch_axis >= 0

        self.ch_axis = ch_axis
        if hasattr(self.activation_post_process, "ch_axis"):
            self.activation_post_process.ch_axis = self.ch_axis

        Method that sets the channel axis. This method should only be called before starting training.

        Args:
            ch_axis (:obj:`int`): Non-negative channel axis index.
        
coremltools/optimize/torch/quantization/modules/learnable_fake_quantize.py