Fix FQ mask in 8da4w QAT (#199)

Co-authored-by: Jerry Zhang <jerryzh168@gmail.com>
pytorch · May 3, 2024 · be30a7f · be30a7f
1 parent 5364de6
commit be30a7f
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