diff --git a/distiller/quantization/range_linear.py b/distiller/quantization/range_linear.py index 5a7c89bbd1b5c0a4dcc8b041c1858fa90b3c811d..3360ed564a832d1069fe619afee16d7e45814110 100644 --- a/distiller/quantization/range_linear.py +++ b/distiller/quantization/range_linear.py @@ -371,8 +371,10 @@ class RangeLinearQuantWrapper(nn.Module): self.register_buffer('num_forwards', torch.zeros(1, dtype=torch.long)) def named_acts_quant_params(self): - yield 'output_scale', self.output_scale - yield 'output_zero_point', self.output_zero_point + if self.preset_act_stats: + # Output scale buffers are saved in the model only when stats are used + yield 'output_scale', self.output_scale + yield 'output_zero_point', self.output_zero_point def forward(self, *inputs): if self.training: