Applying constraint on (Conv+BN+ReLU) folded weights

Hi,
I am trying to add some constraints to the weights and activations in quantization aware training. E.g. I want to make them power of 2. That is weights/activations can take values of 0.125, 0.25, 0.5, 1, 2, 4, 8, e.t.c. I changed the range calculation to achieve this constraints.

It works well for standalone Conv or FC layers. Figure below shows an example: of a conv layer with scale of 128. Mutiplying each weight with 128 will give an integer.

However, for BN followed by conv , our constraints are no longer working . Since, it is ensured that fakequant layer is not inserted between Conv and BN, I was expecting the folded weights to be quantized instead of the standalone conv weights. Now, the weights looks as shown:

Contrary to the expectation of the folded weights being quantized, it looks like that the conv layer is quantized independently and the BN layer is not quantized at all. Finally, during tflite conversion, when the two layers are fused, we don't get the expected "power of 2" weights.

I would like to know what has to be done to make sure that the folded weights (conv+BN) are quantized instead of the standalone conv weights? My ultimate goal is to make the folded weight as power of 2.

System information Linux

TensorFlow version (installed from source or binary): 2.2, Binary

TensorFlow Model Optimization version (installed from source or binary): Binary

Python version: 3.7