captures graph tensor AutoCastDistributedVariable

[sky2code]

System information

• OS Platform and Distribution (e.g., win11):
• TensorFlow version and how it was installed (tf==2.8):
• TensorFlow-Addons version and how it was installed (tfa==0.16.1):
• Python version:(3.9.7)
• Is GPU used? (yes):

Describe the bug PS D:\stock\source> python .\main_retrain.py Number of devices: 1 Epoch 1: val_loss improved from inf to 199.99988, saving model to ..\results\best\CheckPoint Traceback (most recent call last): File "D:\stock\source\main_retrain.py", line 116, in history = model.fit(train, File "C:\Users\zhoug\anaconda3\lib\site-packages\keras\utils\traceback_utils.py", line 67, in error_handler raise e.with_traceback(filtered_tb) from None File "C:\Users\zhoug\anaconda3\lib\site-packages\tensorflow_addons\callbacks\average_model_checkpoint.py", line 86, in _save_model return super()._save_model(*args, **kwargs) ValueError: Unable to save function b'__inference_transformer_layer_call_fn_4234' because it captures graph tensor AutoCastDistributedVariable:{ 0: <tf.Variable 'kernel:0' shape=(512, 8, 512) dtype=float32> } from a parent function which cannot be converted to a constant with tf.get_static_value. A clear and concise description of what the bug is.

Code to reproduce the issue

Provide a reproducible test case that is the bare minimum necessary to generate the problem.

Other info / logs

Include any logs or source code that would be helpful to diagnose the problem. If including tracebacks, please include the full traceback. Large logs and files should be attached.

[yang-stressfree]

I got similar error, details and code to reproduce here https://github.com/tensorflow/addons/issues/2675

[seanpmorgan]

See other issue.