decorate fake quant with torch.no_grad (#8)

gradients shouldn't be computed for Q/DQ in QAT
neuralmagic · Apr 15, 2024 · 0ac6f60 · 0ac6f60
1 parent 129770e
commit 0ac6f60
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/src/sparsetensors/quantization/lifecycle/forward.py b/src/sparsetensors/quantization/lifecycle/forward.py
@@ -24,6 +24,7 @@
 __all__ = ["wrap_module_forward_quantized"]
 
 
+@torch.no_grad()
 def quantize(
     x: torch.Tensor,
     scale: torch.Tensor,
@@ -39,6 +40,7 @@ def quantize(
     )
 
 
+@torch.no_grad()
 def dequantize(
     x_q: torch.Tensor,
     scale: torch.Tensor,
@@ -47,6 +49,7 @@ def dequantize(
     return (x_q - zero_point) * scale
 
 
+@torch.no_grad()
 def fake_quantize(
     x: torch.Tensor,
     scale: torch.Tensor,