Fix FQ mask in 8da4w QAT

pytorch · May 1, 2024 · fec8510 · fec8510
1 parent 6ae2c0b
commit fec8510
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/torchao/quantization/prototype/qat.py b/torchao/quantization/prototype/qat.py
@@ -183,7 +183,7 @@ def forward(ctx, input, scales, zero_points, quant_min, quant_max):
         q = input.div(scales).add(zero_points).round()
         dq = q.clamp(quant_min, quant_max).sub(zero_points).mul(scales)
         # TODO: do we need this mask?
-        mask = torch.logical_and((q >= quant_min), (dq <= quant_max))
+        mask = torch.logical_and((q >= quant_min), (q <= quant_max))
         ctx.save_for_backward(mask)
         return dq