Fix (nn): correct bias handlind in ScaleBias

Xilinx · May 3, 2023 · acb5eda · acb5eda
1 parent 7f83831
commit acb5eda
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/src/brevitas/nn/quant_scale_bias.py b/src/brevitas/nn/quant_scale_bias.py
@@ -31,7 +31,10 @@ def __init__(self, num_features: int, bias: bool, runtime_shape=(1, -1, 1, 1)):
         self.runtime_shape = runtime_shape
 
     def forward(self, input):
-        return input * self.weight.view(self.runtime_shape) + self.bias.view(self.runtime_shape)
+        out = input * self.weight.view(self.runtime_shape)
+        if self.bias:
+            out += self.bias.view(self.runtime_shape)
+        return out
 
 
 class QuantScaleBias(QuantWBIOL, ScaleBias):