Unknown layer: KerasLayer. Please ensure this object is passed to the `custom_objects` argument. Error when trying to do Quantization aware training

I have done simple transfer learning for a custom dataset using a mobilenetv2 model available in Tensorflow Hub. However I'm getting this error when I try to do Quantization aware training. I don't understand how to solve it exactly. What should I be doing? Below is my code and the error I'm getting

import numpy as np
import time

import PIL.Image as Image
import matplotlib.pylab as plt

import tensorflow as tf
from tensorflow import keras
from keras.preprocessing.image import ImageDataGenerator

import tensorflow_hub as hub

import datetime


from google.colab import drive
drive.mount('/content/gdrive', force_remount=True)
dataset_path = '/content/gdrive/MyDrive/images/'

image_size = (224,224)
batch_size = 10

train_datagen = ImageDataGenerator(rescale=1./255, 
                                   rotation_range=15,
                                   zoom_range = (0.95,0.95),
                                   width_shift_range=0.1,
                                   height_shift_range=0.1,
                                   validation_split=0.2,
                                   dtype = tf.float32,
                                  )

validation_datagen = ImageDataGenerator(rescale=1./255, validation_split=0.2, dtype = tf.float32,)

train_batches = train_datagen.flow_from_directory(
    dataset_path,
    target_size = image_size,
    batch_size = batch_size,
    color_mode='rgb',
    class_mode = 'categorical',
    shuffle = True,
    seed = 123,
    subset = 'training', )

validation_batches = validation_datagen.flow_from_directory(
    dataset_path,
    target_size = image_size,
    batch_size = batch_size,
    color_mode='rgb',
    class_mode = 'categorical',
    shuffle = True,
    seed = 123,
    subset = 'validation', )

test_batches = validation_datagen.flow_from_directory(
    dataset_path,
    target_size = image_size,
    batch_size = batch_size,
    color_mode='rgb',
    class_mode = 'categorical',
    shuffle = True,
    seed = 123,
    subset = 'validation', )

from keras import layers
import tensorflow_hub as hub


url = "https://tfhub.dev/google/experts/bit/r50x1/in21k/bird/1"
classifier_model = url
base_model = hub.KerasLayer(classifier_model, input_shape=image_size+(3,), trainable=False)

num_of_birds = 10

model = tf.keras.Sequential([
    base_model,
    tf.keras.layers.Dense(num_of_birds, activation='softmax'),
])

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

model.fit_generator(train_batches, epochs=5)

model.evaluate(test_batches)

import tensorflow_model_optimization as tfmot

quantize_model = tfmot.quantization.keras.quantize_model

q_aware_model = quantize_model(model)

I get this error right when I try to run the above line of code

ValueError                                Traceback (most recent call last)
/usr/local/lib/python3.7/dist-packages/tensorflow_model_optimization/python/core/quantization/keras/quantize.py in quantize_apply(model, scheme)
    442   try:
--> 443     model_copy = _clone_model_with_weights(model)
    444   except ValueError as er:

12 frames
ValueError: Unknown layer: KerasLayer. 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 in quantize_apply(model, scheme)
        447         'Keras layers or objects in your model. Please specify them via '
        448         '`quantize_scope` for your calls to `quantize_model` and '
    --> 449         '`quantize_apply`. [%s].' % er)
        450 
        451   # 2. Remove QuantizeAnnotate wrappers from the layers in the model. This
    
    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`. [Unknown layer: KerasLayer. 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.].