Update pre_scaling.py

Xilinx · Nov 2, 2023 · b5f0e9e · b5f0e9e
1 parent 6cdadcd
commit b5f0e9e
Showing 1 changed file with 1 addition and 2 deletions.
diff --git a/src/brevitas/core/scaling/pre_scaling.py b/src/brevitas/core/scaling/pre_scaling.py
@@ -15,8 +15,7 @@
 from brevitas.core.stats.stats_wrapper import _Stats
 from brevitas.function import abs_binary_sign_grad
 
-__all__ = [
-    "ParameterPreScalingWeightNorm",]
+__all__ = ["ParameterPreScalingWeightNorm", "AccumulatorAwareParameterPreScaling"]
 
 
 class ParameterPreScalingWeightNorm(brevitas.jit.ScriptModule):