use keras.Sequential(Conv2D,BatchNormalization) convert quantize_model got error

[zhen8838]

Describe the bug

convert Sequential model to quantize_model will got error, But use Functional model will be success.

System information

TensorFlow installed from (source or binary): binary

TensorFlow version: tf 2.1

TensorFlow Model Optimization version: 0.3.0

Python version: 3.7.4

Describe the expected behavior

Describe the current behavior

model = keras.Sequential([
    keras.layers.InputLayer(input_shape=(28, 28)),
    keras.layers.Reshape(target_shape=(28, 28, 1)),
    keras.layers.Conv2D(filters=12, kernel_size=(3, 3)),
    keras.layers.BatchNormalization(),
    keras.layers.ReLU(),
    keras.layers.MaxPooling2D(pool_size=(2, 2)),
    keras.layers.Flatten(),
    keras.layers.Dense(10, activation=tf.nn.softmax)
])

quantize_model = tfmot.quantization.keras.quantize_model

# q_aware stands for for quantization aware.
q_aware_model = quantize_model(model)

ValueError: Input 0 of layer conv2d is incompatible with the layer: expected ndim=4, found ndim=2. Full shape received: [None, 196]

Code to reproduce the issue
link: https://colab.research.google.com/drive/1GZzAhUTgcDUFGQUvJu59be0AgAoBkhrw?usp=sharing

Screenshots
If applicable, add screenshots to help explain your problem.

Additional context
Add any other context about the problem here.

[miaout17]

@alanchiao do you mind to take a look? Thanks!

[alanchiao]

This was fixed in #402

In the meantime, you'll need to build from source to enable it. It'll be a part of the next release after 0.3.0.