Fix LSTM

Xilinx · Oct 1, 2024 · 609a164 · 609a164
1 parent f44fe35
commit 609a164
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/src/brevitas/core/stats/stats_wrapper.py b/src/brevitas/core/stats/stats_wrapper.py
@@ -115,9 +115,11 @@ def __init__(
 
     @brevitas.jit.script_method
     def forward(self, x: Optional[torch.Tensor] = None) -> torch.Tensor:
-        stats_input = self.first_tracked_param(x)
         if self.extra_tracked_params_list is not None:
+            stats_input = self.first_tracked_param(None)
             for extra_tracked_param in self.extra_tracked_params_list:
                 stats_input = extra_tracked_param(stats_input)
+        else:
+            stats_input = self.first_tracked_param(x)
         out = self.stats(stats_input)
         return out