add forward-onnly backends for debugging

mobiusml · Mar 18, 2024 · 4a95c9e · 4a95c9e
1 parent 9e7249e
commit 4a95c9e
Showing 1 changed file with 5 additions and 0 deletions.
diff --git a/hqq/core/quantize.py b/hqq/core/quantize.py
@@ -246,6 +246,10 @@ class HQQBackend(Enum):
     PYTORCH_BACKPROP_COMPILE = "forward_pytorch_backprop_compile"
     ATEN_BACKPROP = "forward_aten_backprop"
 
+    PYTORCH_FORWARD = "forward_pytorch"
+    PYTORCH_FORWARD_COMPILE = "forward_pytorch_compile"
+    ATEN_FORWARD = "forward_aten"
+
 
 # No cache: less memory, slower
 class HQQMatmulNoCacheDeq(torch.autograd.Function):
@@ -779,6 +783,7 @@ def forward_aten(self, x: Tensor) -> Tensor:
         out = torch.matmul(x, W_est.t())
         if self.bias is not None:
             out += self.bias
+
         return out
 
     def forward_aten_backprop(self, x: Tensor) -> Tensor: