Fix for cache + decoupled weight quant

Xilinx · Nov 28, 2024 · 84c4309 · 84c4309
1 parent 1bc7df9
commit 84c4309
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/src/brevitas/proxy/parameter_quant.py b/src/brevitas/proxy/parameter_quant.py
@@ -263,7 +263,7 @@ def forward(
             self._cached_act = cached_inp
 
         if self.is_quant_enabled:
-            if quant_input is None:
+            if quant_input is None or isinstance(quant_input, Tensor):
                 assert self._cached_act is not None, "No cached quant input found. Enable caching and perform a forward pass"
                 quant_input = self._cached_act
             else: