def set_device(self, device: _torch.device):
        """
        Method that sets the device of the parameters.

        Args:
            device (:py:class:`torch.device`): Device where to hold learnable parameters.
        """

        self.to(device)
        self.device = device

        Method that sets the device of the parameters.

        Args:
            device (:py:class:`torch.device`): Device where to hold learnable parameters.
        
coremltools/optimize/torch/quantization/modules/learnable_fake_quantize.py