def _setup_fake_quant_module_device(
        self, model: _fx.GraphModule, example_inputs: _Tuple[_Any, ...]
    ):
        """
        Set device for all fake quantize modules by inferring from model and/or data
        """
        # Record the device of the model
        count_params = _count_model_params(model)
        if count_params > 0:
            self._device = next(model.parameters()).device
        elif len(example_inputs) > 0:
            self._device = example_inputs[0].device
        else:
            self._device = _torch.device("cpu")

        for name, module in model.named_modules(remove_duplicate=True):
            # Set channel axis and param shape for weight observers
            if (
                hasattr(module, "weight_fake_quant")
                and module.weight_fake_quant is not None
                and hasattr(module.weight_fake_quant, "qscheme")
            ):
                if _is_per_channel_quant(module.weight_fake_quant.qscheme):
                    if hasattr(module.weight_fake_quant, "set_ch_axis"):
                        weight_ch_axis = module.weight_fake_quant.ch_axis
                        module.weight_fake_quant.set_ch_axis(weight_ch_axis)
                    if hasattr(module.weight_fake_quant, "set_param_shape"):
                        weight_shape = module.weight.shape
                        module.weight_fake_quant.set_param_shape(weight_shape)

            # Set device for weight and activation observers
            if (
                hasattr(module, "weight_fake_quant")
                and module.weight_fake_quant is not None
                and hasattr(module.weight_fake_quant, "set_device")
            ):
                module.weight_fake_quant.set_device(self._device)
            elif not name.endswith(".weight_fake_quant") and hasattr(module, "set_device"):
                module.set_device(self._device)

        Set device for all fake quantize modules by inferring from model and/or data
        
coremltools/optimize/torch/quantization/_configure.py