Problem regarding loading a QAT SavedModel

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
Currently, I am using QAT inside a custom training loop context and it goes pretty smoothly. But the problem starts surfacing when I try to load the serialized SavedModel (Colab Notebook) -

KeyError: '__inference_expanded_conv_depthwise_layer_call_and_return_conditional_losses_330835'

System information

I am using Colab with TensorFlow 2.4.1.

Describe the expected behavior

The SavedModel should load as expected.

Describe the current behavior

The SavedModel is not getting loaded as expected and the predictions of the QAT model are extremely spurious.

Code to reproduce the issue
This zip file contains a Jupyter Notebook. Although it uses 4 GPUs to distribute the training, I believe if you disable the strategy scope it will run just fine on a Colab Notebook.

Knowledge_Distillation_MobileNetV2_QAT.ipynb.zip

Screenshots
If applicable, add screenshots to help explain your problem.

Additional context
Add any other context about the problem here.