Feat (mse): use grid search for scale

Xilinx · Dec 3, 2024 · 0682b94 · 0682b94
1 parent d4e001e
commit 0682b94
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/src/brevitas/quant/base.py b/src/brevitas/quant/base.py
@@ -498,6 +498,7 @@ class MSESymmetricScaleSubInjector(ExtendedInjector):
     type = (this << 1).type
     permute_dims = (this << 1).permute_dims
     inner_stats_input_view_shape_impl = (this << 1).inner_stats_input_view_shape_impl
+    mse_search_method = 'grid'
 
 
 class MSEAsymmetricScaleSubInjector(ExtendedInjector):
@@ -510,6 +511,7 @@ class MSEAsymmetricScaleSubInjector(ExtendedInjector):
     dtype = (this << 1).dtype
     permute_dims = (this << 1).permute_dims
     inner_stats_input_view_shape_impl = (this << 1).inner_stats_input_view_shape_impl
+    mse_search_method = 'grid'
 
 
 class MSEZeroPointSubInjector(ExtendedInjector):