Quantizing model at runtime during training results in non-learned quantization

Prior to filing: check that this should be a bug instead of a feature request. Everything supported, including the compatible versions of TensorFlow, is listed in the overview page of each technique. For example, the overview page of quantization-aware training is here. An issue for anything not supported should be a feature request.

Describe the bug
I want to perform quantization-aware training. I tried to start a training with a non-quantized model and to apply quantization at a certain epoch. The training goes well with a good validation accuracy. However, when I load the trained model to perform some testing, it appears that the quantization parameters have not been learned and the quantization is not properly applied to the input, resulting in black images and very bad results.

System information

TensorFlow version (installed from source or binary): 2.4.1 binary

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

Python version: 3.6.9

Describe the expected behavior
If the quantization is applied at runtime, the quantization parameters should be learned as if the model was quantized from the start of the training

Describe the current behavior
When the quantization is applied at runtime, the training goes well but the saved model does not quantize the input properly

Code to reproduce the issue
Relevant code attached.

@tf.function
def train_step(x, y):
	with tf.GradientTape() as tape:
		probs = model(x, training=True)
		loss_value = loss_fn(y, probs)
		
	grads = tape.gradient(loss_value, model.trainable_weights)
	optimizer.apply_gradients(zip(grads, model.trainable_weights))
	train_acc_metric.update_state(y, probs)
	return loss_value, probs

def custom_quantization(layer):
	return tfmot.quantization.keras.quantize_annotate_layer(layer)


# load datasets
train_dataset, val_dataset = get_datasets()

# define new model or load saved model
model = get_model()
loss_fn = tf.keras.losses.CategoricalCrossentropy()
optimizer = tf.keras.optimizers.Adam(learning_rate=0.001)

for epoch in range(1, epochs):
        if epoch == quantization_epoch_start:
		print("Quantizing model...")
		annotated_model = tf.keras.models.clone_model(model, clone_function=custom_quantization)

		with tfmot.quantization.keras.quantize_scope():
			model = tfmot.quantization.keras.quantize_apply(annotated_model)

	for step, (x_batch_train, y_batch_train) in enumerate(train_dataset):		
                total_step += 1
		loss_value, scores = train_step(x_batch_train, y_batch_train)

        if total_step % 2000 == 0:
	        print("saving model...")
        	model.save(os.path.join(hparam_ckpt_path, "step-%d.h5" % (total_step)), save_format='h5')