Fix (quant): propagate device and dtype in subinjector (#942)

Xilinx · Apr 29, 2024 · cae4004 · cae4004
1 parent bae1e26
commit cae4004
Showing 1 changed file with 6 additions and 0 deletions.
diff --git a/src/brevitas/quant/base.py b/src/brevitas/quant/base.py
@@ -435,6 +435,8 @@ class MSESymmetricScaleSubInjector(MSESubInjectorBase):
  mse_init_op = AbsMax
  stats_impl = MSE
  stats_reduce_dim = (this << 1).stats_reduce_dim
+ device = (this << 1).device
+ type = (this << 1).type
 
 
 class MSEAsymmetricScaleSubInjector(MSESubInjectorBase):
@@ -443,6 +445,8 @@ class MSEAsymmetricScaleSubInjector(MSESubInjectorBase):
  mse_init_op = AbsMinMax
  stats_impl = MSE
  stats_reduce_dim = (this << 1).stats_reduce_dim
+ device = (this << 1).device
+ type = (this << 1).type
 
 
 class MSEZeroPointSubInjector(MSESubInjectorBase):
@@ -453,6 +457,8 @@ class MSEZeroPointSubInjector(MSESubInjectorBase):
  mse_search_method = 'grid'
  stats_impl = MSE
  stats_reduce_dim = (this << 1).stats_reduce_dim
+ device = (this << 1).device
+ type = (this << 1).type
 
 
 class MSEAsymmetricScale(ExtendedInjector):