Hi @XiaoquanKong,
I run the old code (set the dtype = tf.int32) several times on my PC and all got the same error messages. I tried to run it on Google Colab with and without the GPU and they all get the same errors in TF2.6.0.(But I couldn’t run the GPU version properly on Google Colab even I set the dtype = tf.float32 because it cannot find the cudnn module)  Please find My code and the testing data through the links and let me know if they cannot work properly.
About the other related question, I run the same code with the same testing data provided above on Google Colab, and it only shows a few minor warnings when saving:
WARNING:absl:Found untraced functions such as dense_1_layer_call_and_return_conditional_losses, dense_1_layer_call_fn, dense_1_layer_call_fn, dense_1_layer_call_and_return_conditional_losses, dense_1_layer_call_and_return_conditional_losses while saving (showing 5 of 15). These functions will not be directly callable after loading.

I am wondering if you ever test the code with tensorflow GPU version (i.e. tensorflow-gpu=2.6.0) because this is the main difference between my tests running on Google Colab (tensorflow=2.6.0) and on my PC (tensorflow-gpu=2.6.0). So far I couldn’t run the code with tensorflow GPU version on Google Colab and will let you know once I fix it.
Thank you very much.