fix for multiple inits

Xilinx · Jan 6, 2025 · b45f298 · b45f298
1 parent 9146d5b
commit b45f298
Showing 1 changed file with 4 additions and 3 deletions.
diff --git a/src/brevitas/core/scaling/standalone.py b/src/brevitas/core/scaling/standalone.py
@@ -376,9 +376,10 @@ def __init__(
         self.restrict_threshold_pre = restrict_threshold_impl.restrict_init_module()
 
     def init_scale(self):
-        self.restrict_inplace_preprocess(self.buffer)
-        inplace_tensor_mul(self.value.detach(), self.buffer)
-        self.counter = self.counter + 1
+        if self.counter <= self.collect_stats_steps:
+            self.restrict_inplace_preprocess(self.buffer)
+            inplace_tensor_mul(self.value.detach(), self.buffer)
+            self.counter = self.collect_stats_steps + 1
 
     @brevitas.jit.script_method
     def training_forward(self, stats_input: Tensor, threshold: Tensor) -> Tensor: