Skip to content

Commit 90f5ccd

Browse files
ratgrtensorflower-gardener
authored andcommitted
Fix an issue where multiple instances of DefaultNBitQuantizeRegistry could interfere with each other's configuration.
PiperOrigin-RevId: 931228650
1 parent 97811c3 commit 90f5ccd

1 file changed

Lines changed: 9 additions & 3 deletions

File tree

‎tensorflow_model_optimization/python/core/quantization/keras/experimental/default_n_bit/default_n_bit_quantize_registry.py‎

Lines changed: 9 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -185,9 +185,15 @@ def __init__(self, disable_per_axis=False,
185185
self._num_bits_activation = num_bits_activation
186186
self._layer_quantize_map = {}
187187
for quantize_info in self._LAYER_QUANTIZE_INFO:
188-
quantize_info.num_bits_weight = num_bits_weight
189-
quantize_info.num_bits_activation = num_bits_activation
190-
self._layer_quantize_map[quantize_info.layer_type] = quantize_info
188+
new_quantize_info = _QuantizeInfo(
189+
layer_type=quantize_info.layer_type,
190+
weight_attrs=quantize_info.weight_attrs,
191+
activation_attrs=quantize_info.activation_attrs,
192+
quantize_output=quantize_info.quantize_output,
193+
num_bits_weight=num_bits_weight,
194+
num_bits_activation=num_bits_activation,
195+
)
196+
self._layer_quantize_map[new_quantize_info.layer_type] = new_quantize_info
191197

192198
# Hack for `Activation` layer. That is the only layer with a separate
193199
# QuantizeConfig.

0 commit comments

Comments
 (0)