diff --git a/concrete/quantization/post_training.py b/concrete/quantization/post_training.py index f60fb8614..4736fcb1b 100644 --- a/concrete/quantization/post_training.py +++ b/concrete/quantization/post_training.py @@ -88,7 +88,12 @@ class PostTrainingAffineQuantization: # Create a QuantizedLinear layer q_weights = self.quant_params[f"{name}.weight"] q_bias = self.quant_params[f"{name}.bias"] - q_layer = QuantizedLinear(self.n_bits, q_weights, q_bias) + # Check if layer is last layer from the model + if name == list(self.numpy_model.torch_model.named_children())[-1][0]: + # If last layer, we can use 7 bits (maximum allowed) of precision. + q_layer = QuantizedLinear(7, q_weights, q_bias) + else: + q_layer = QuantizedLinear(self.n_bits, q_weights, q_bias) # Calibrate and get new calibration_data for next layer/activation calibration_data = self._calibrate_layers_activation( name, q_layer, calibration_data