QAT model saved not working correctly, and cause errors in tf2onnx

[QW-Is-Here]

Prerequisites

Please answer the following questions for yourself before submitting an issue.

• [x] I am using the latest TensorFlow Model Garden release and TensorFlow 2.4
• [ ] I am reporting the issue to the correct repository. (Model Garden official or research directory)
• [ ] I checked to make sure that this issue has not been filed already.

1. The entire URL of the file you are using

https://github.com/tensorflow/models/tree/master/official/...

2. Describe the bug

After quantizing the model, the model cannot save correctly, and can not build an onnx model from tf.keras QAT saved model

3. Steps to reproduce

Same as this script https://colab.research.google.com/drive/1LuNdGlQ2XdlCNyua_araGr6RYrwhiHPm?usp=sharing

and after saved model when i try to convert saved model to onnx model i got errors.

>>> q_aware_mobilenet.save('test_saved_model_qat/saved_mobilenet/'
... )

WARNING:absl:Found untraced functions such as conv1_layer_call_fn, conv1_layer_call_and_return_conditional_losses, conv1_relu_layer_call_fn, conv1_relu_layer_call_and_return_conditional_losses, conv_dw_1_relu_layer_call_fn while saving (showing 5 of 205). These functions will not be directly callable after loading.
WARNING:absl:Found untraced functions such as conv1_layer_call_fn, conv1_layer_call_and_return_conditional_losses, conv1_relu_layer_call_fn, conv1_relu_layer_call_and_return_conditional_losses, conv_dw_1_relu_layer_call_fn while saving (showing 5 of 205). These functions will not be directly callable after loading.
INFO:tensorflow:Assets written to: test_saved_model_qat/saved_mobilenet/assets
INFO:tensorflow:Assets written to: test_saved_model_qat/saved_mobilenet/assets

python -m tf2onnx.convert --saved-model=saved_mobilenet --opset=13 --output=model.onnx  --tag=serve

ce:GPU:1 with 22033 MB memory) -> physical GPU (device: 1, name: NVIDIA GeForce RTX 3090, pci bus id: 0000:4a:00.0, compute capability: 8.6)
Traceback (most recent call last):
  File "/usr/lib/python3.6/runpy.py", line 193, in _run_module_as_main
    "__main__", mod_spec)
  File "/usr/lib/python3.6/runpy.py", line 85, in _run_code
    exec(code, run_globals)
  File "/usr/local/lib/python3.6/dist-packages/tf2onnx/convert.py", line 706, in <module>
    main()
  File "/usr/local/lib/python3.6/dist-packages/tf2onnx/convert.py", line 241, in main
    use_graph_names=args.use_graph_names)
  File "/usr/local/lib/python3.6/dist-packages/tf2onnx/tf_loader.py", line 615, in from_saved_model
    tag, signatures, concrete_function, large_model, use_graph_names)
  File "/usr/local/lib/python3.6/dist-packages/tf2onnx/tf_loader.py", line 552, in _from_saved_model_v2
    imported = tf.saved_model.load(model_path, tags=tag)  # pylint: disable=no-value-for-parameter
  File "/usr/local/lib/python3.6/dist-packages/tensorflow/python/saved_model/load.py", line 859, in load
    return load_internal(export_dir, tags, options)["root"]
  File "/usr/local/lib/python3.6/dist-packages/tensorflow/python/saved_model/load.py", line 890, in load_internal
    ckpt_options, filters)
  File "/usr/local/lib/python3.6/dist-packages/tensorflow/python/saved_model/load.py", line 160, in __init__
    self._load_all()
  File "/usr/local/lib/python3.6/dist-packages/tensorflow/python/saved_model/load.py", line 256, in _load_all
    self._load_nodes()
  File "/usr/local/lib/python3.6/dist-packages/tensorflow/python/saved_model/load.py", line 434, in _load_nodes
    node, setter = self._recreate(proto, node_id)
  File "/usr/local/lib/python3.6/dist-packages/tensorflow/python/saved_model/load.py", line 552, in _recreate
    return factory[kind]()
  File "/usr/local/lib/python3.6/dist-packages/tensorflow/python/saved_model/load.py", line 541, in <lambda>
    "function": lambda: self._recreate_function(proto.function),
  File "/usr/local/lib/python3.6/dist-packages/tensorflow/python/saved_model/load.py", line 580, in _recreate_function
    proto, self._concrete_functions), setattr
  File "/usr/local/lib/python3.6/dist-packages/tensorflow/python/saved_model/function_deserialization.py", line 277, in recreate_function
    concrete_function_objects.append(concrete_functions[concrete_function_name])
KeyError: '__inference_conv_dw_1_layer_call_fn_240991'

2. TensorFlow 2.0 python -c "import tensorflow as tf; print(tf.version.GIT_VERSION, tf.version.VERSION)" -->

2022-08-02 07:45:22.944743: I tensorflow/stream_executor/platform/default/dso_loader.cc:49] Successfully opened dynamic library libcudart.so.11.0
v2.4.0-0-g582c8d236cb 2.4.0