Support Dense+BatchNorm in QAT

Hi Everyone,
i built an NN with BatchNormalization layer and i have tried to quantize the whole model for EdgeTPU application. I have read that i can use this layer after Dense or Conv2D layer in this overview: https://www.tensorflow.org/model_optimization/guide/quantization/training#api_compatibility. But it doesn't work.

Tensorflow: 2.1
Tensorflow Model Optimization : 0.3.0
Python : 3.7.7

the packages have been installed with pip.

This is my code :

def create_model():
        

    model = tf.keras.Sequential([
        tf.keras.layers.Dense( units=5,input_shape=(224,224,3)),
        tf.keras.layers.BatchNormalization(),
        tf.keras.layers.MaxPool2D(),    
        tf.keras.layers.Dense( units=5),        
        tf.keras.layers.MaxPool2D(),      
        tf.keras.layers.Dropout(rate = 0.4),       
        tf.keras.layers.Conv2D(kernel_size = 4 , activation="softmax", filters = 32,kernel_initializer='he_normal'),
        tf.keras.layers.MaxPool2D(),
        tf.keras.layers.Dropout(rate = 0.4),
        tf.keras.layers.Conv2D(kernel_size = 3 , activation="softmax", filters = 6,kernel_initializer='he_normal'),
        tf.keras.layers.GlobalMaxPool2D()
    ])

    model.compile(optimizer="adam",loss="categorical_crossentropy",metrics=['accuracy'])

    return model

model = create_model()
quant_aware_model = tfmot.quantization.keras.quantize_model(model)
quant_aware_model.summary()

converter = tf.lite.TFLiteConverter.from_keras_model(quant_aware_model)
converter.optimizations = [tf.lite.Optimize.DEFAULT]

quantized_tflite_model = converter.convert()

and this is the error

RuntimeError: Layer batch_normalization:<class 'tensorflow.python.keras.layers.normalization_v2.BatchNormalization'> is not supported. You can quantize 
this layer by passing a `tfmot.quantization.keras.QuantizeConfig` instance to the `quantize_annotate_layer` API.

Could i fix it?