Quantize SSD models with TF TensorFlow model optimization tool

I'm trying to compress the ssd_mobilenet_v1_fpn_640x640_coco17_tpu-8 model (from model zoo) with tensorflow optimization tool, or more specifically, tensorflow_model_optimization, which supports quantizing TF/Keras models by choosing which to quantize.

import tensorflow_model_optimization as tfmot
model = keras.models.load_model(model_dir)
quantize_model = tfmot.quantization.keras.quantize_model
q_aware_model = quantize_model(model)

However, I had this warning and error:

WARNING:tensorflow:SavedModel saved prior to TF 2.4 detected when loading Keras model. Please ensure that you are saving the model with model.save() or tf.keras.models.save_model(), *NOT* tf.saved_model.save(). To confirm, there should be a file named "keras_metadata.pb" in the SavedModel directory.
---------------------------------------------------------------------------
AttributeError                            Traceback (most recent call last)
<ipython-input-34-189333fbe081> in <module>
      1 tf.keras.backend.clear_session()
----> 2 model = keras.models.load_model(model_dir) # version not matched
      3 
      4 configs = config_util.get_configs_from_pipeline_file(config_fn)
      5 model_config = configs['model']

~/anaconda3/lib/python3.7/site-packages/tensorflow/python/keras/saving/save.py in load_model(filepath, custom_objects, compile, options)
    211       if isinstance(filepath, six.string_types):
    212         loader_impl.parse_saved_model(filepath)
--> 213         return saved_model_load.load(filepath, compile, options)
    214 
    215   raise IOError(

~/anaconda3/lib/python3.7/site-packages/tensorflow/python/keras/saving/saved_model/load.py in load(path, compile, options)
    156 
    157   # Finalize the loaded layers and remove the extra tracked dependencies.
--> 158   keras_loader.finalize_objects()
    159   keras_loader.del_tracking()
    160 

~/anaconda3/lib/python3.7/site-packages/tensorflow/python/keras/saving/saved_model/load.py in finalize_objects(self)
    610         layers_revived_from_config.append(node)
    611 
--> 612     _finalize_saved_model_layers(layers_revived_from_saved_model)
    613     _finalize_config_layers(layers_revived_from_config)
    614 

~/anaconda3/lib/python3.7/site-packages/tensorflow/python/keras/saving/saved_model/load.py in _finalize_saved_model_layers(layers)
    799         call_fn = _get_keras_attr(layer).call_and_return_conditional_losses
    800         if call_fn.input_signature is None:
--> 801           inputs = infer_inputs_from_restored_call_function(call_fn)
    802         else:
    803           inputs = call_fn.input_signature[0]

~/anaconda3/lib/python3.7/site-packages/tensorflow/python/keras/saving/saved_model/load.py in infer_inputs_from_restored_call_function(fn)
   1110   for concrete in fn.concrete_functions[1:]:
   1111     spec2 = concrete.structured_input_signature[0][0]
-> 1112     spec = nest.map_structure(common_spec, spec, spec2)
   1113   return spec
   1114 

~/anaconda3/lib/python3.7/site-packages/tensorflow/python/util/nest.py in map_structure(func, *structure, **kwargs)
    887 
    888   return pack_sequence_as(
--> 889       structure[0], [func(*x) for x in entries],
    890       expand_composites=expand_composites)
    891 

~/anaconda3/lib/python3.7/site-packages/tensorflow/python/util/nest.py in <listcomp>(.0)
    887 
    888   return pack_sequence_as(
--> 889       structure[0], [func(*x) for x in entries],
    890       expand_composites=expand_composites)
    891 

~/anaconda3/lib/python3.7/site-packages/tensorflow/python/keras/saving/saved_model/load.py in common_spec(x, y)
   1102   """
   1103   def common_spec(x, y):
-> 1104     common_shape = get_common_shape(x.shape, y.shape)
   1105     if isinstance(x, sparse_tensor.SparseTensorSpec):
   1106       return sparse_tensor.SparseTensorSpec(common_shape, x.dtype)

AttributeError: 'str' object has no attribute 'shape'

The inputted SavedModel is converted by object_detection/export_tflite_graph_tf2.py. Anyone know how can I load the OD model into a keras model?