Fix

Xilinx · Dec 5, 2024 · 260009d · 260009d
1 parent d5bf3d3
commit 260009d
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 2 deletions.
diff --git a/src/brevitas/core/scaling/runtime.py b/src/brevitas/core/scaling/runtime.py
@@ -30,6 +30,7 @@ def __init__(
             scaling_stats_input_concat_dim: int,
             tracked_parameter_list: List[torch.nn.Parameter],
             scaling_shape: Tuple[int, ...],
+            force_parameter: bool = False,
             restrict_scaling_impl: Module = FloatRestrictValue(),
             restrict_threshold_impl: Optional[Module] = None,
             affine_rescaling: bool = False,
@@ -48,7 +49,8 @@ def __init__(
             scaling_shape,
             scaling_stats_input_view_shape_impl,
             scaling_stats_input_concat_dim,
-            tracked_parameter_list)
+            tracked_parameter_list,
+            force_parameter)
         self.stats_scaling_impl = _StatsScaling(
             restrict_scaling_impl,
             restrict_threshold_impl,

diff --git a/src/brevitas/core/scaling/standalone.py b/src/brevitas/core/scaling/standalone.py
@@ -207,6 +207,7 @@ def __init__(
             scaling_stats_input_concat_dim: int,
             tracked_parameter_list: List[torch.nn.Parameter],
             scaling_shape: Tuple[int, ...],
+            force_parameter: bool = False,
             restrict_scaling_impl: Module = FloatRestrictValue(),
             restrict_threshold_impl: Optional[Module] = None,
             scaling_min_val: Optional[float] = None,
@@ -218,7 +219,8 @@ def __init__(
             scaling_shape,
             scaling_stats_input_view_shape_impl,
             scaling_stats_input_concat_dim,
-            tracked_parameter_list)
+            tracked_parameter_list,
+            force_parameter)
 
         # Ensure retro-compatibility with shared threshold/scaling restrict
         if restrict_threshold_impl is None: