Quantization aware training for a Transfer Learning MobileNet model

Hello, I have a MobileNetV2 That I am trying to use for image classification by means of transfer learning, although apparently seems to not work. Initially, I perform transfer learning on my model as follows:

base_model = tf.keras.applications.MobileNetV2(include_top=False)
base_model.trainable = True

inputs = tf.keras.layers.Input(shape=(384, 288, 3), name="input_layer")
x = tf.keras.layers.experimental.preprocessing.Rescaling(1./255)(inputs)
x = base_model(inputs)
x = tf.keras.layers.GlobalAveragePooling2D(name="global_average_pooling_layer")(x)
outputs = tf.keras.layers.Dense(5, activation="softmax", name="output_layer")(x)


model = tf.keras.Model(inputs, outputs)
model.compile(loss='categorical_crossentropy',
              optimizer=tf.keras.optimizers.Adam(),
              metrics=["accuracy"])

history = model.fit(training_dataset,
                                 epochs=30,
                                 steps_per_epoch=len(training_dataset),
                                 validation_data=validation_data,
                                 validation_steps=int(len(validation_data)))

I then followed the steps here to perform model quantization: https://colab.research.google.com/github/tensorflow/model-optimization/blob/master/tensorflow_model_optimization/g3doc/guide/quantization/training_example.ipynb#scrollTo=oq6blGjgFDCW

As the above, I attempted a quantization aware training on my model like this:

model = tf.keras.Model(inputs, outputs)

quantize_model = tfmot.quantization.keras.quantize_model

model = quantize_model(model)
model.compile(loss='categorical_crossentropy',
              optimizer=tf.keras.optimizers.Adam(),
              metrics=["accuracy"])

history = model.fit(training_dataset,
                                 epochs=30,
                                 steps_per_epoch=len(training_dataset),
                                 validation_data=validation_data,
                                 validation_steps=int(len(validation_data)))

It gives me the following error: ValueError: Quantizing a tf.keras Model inside another tf.keras Model is not supported.

Then I tried what was in this link: https://github.com/tensorflow/model-optimization/issues/377#issuecomment-820948555

I tried doing the below attempt:

base_model = tf.keras.applications.MobileNetV2(include_top=False)
base_model.trainable = False
q_base_model = quantize_model(base_model)

inputs = tf.keras.layers.Input(shape=(384, 288, 3), name="input_layer")
x = tf.keras.layers.experimental.preprocessing.Rescaling(1./255)(inputs)
x = base_model(inputs)
x = tf.keras.layers.GlobalAveragePooling2D(name="global_average_pooling_layer")(x)
outputs = tf.keras.layers.Dense(5, activation="softmax", name="output_layer")(x)

model = tf.keras.Model(inputs, outputs)
model = quantize_model(q_base_model)
model.compile(loss='categorical_crossentropy',
              optimizer=tf.keras.optimizers.Adam(),
              metrics=["accuracy"])

history = model.fit(training_dataset,
                                 epochs=30,
                                 steps_per_epoch=len(training_dataset),
                                 validation_data=validation_data,
                                 validation_steps=int(len(validation_data)))

Which outputs the following error:

ValueError                                Traceback (most recent call last)
[/usr/local/lib/python3.7/dist-packages/tensorflow_model_optimization/python/core/quantization/keras/quantize.py](https://localhost:8080/#) in quantize_apply(model, scheme)
    437   try:
--> 438     model_copy = _clone_model_with_weights(model)
    439   except ValueError:

15 frames
ValueError: Unknown layer: QuantizeLayer. Please ensure this object is passed to the `custom_objects` argument. See https://www.tensorflow.org/guide/keras/save_and_serialize#registering_the_custom_object for details.

During handling of the above exception, another exception occurred:

ValueError                                Traceback (most recent call last)
[/usr/local/lib/python3.7/dist-packages/tensorflow_model_optimization/python/core/quantization/keras/quantize.py](https://localhost:8080/#) in quantize_apply(model, scheme)
    439   except ValueError:
    440     raise ValueError(
--> 441         'Unable to clone model. This generally happens if you used custom '
    442         'Keras layers or objects in your model. Please specify them via '
    443         '`quantize_scope` for your calls to `quantize_model` and '

ValueError: Unable to clone model. This generally happens if you used custom Keras layers or objects in your model. Please specify them via quantize_scope for your calls to quantize_model and quantize_apply.
It seems I haven't fully understood how to get quantization aware training done correctly. I'd like to request help on how to properly do QAT on a transfer learning model such as the above example?