issues with @tf.function while training , saving and loading the encoder-decoder model

[mankar1257]

System information

• OS Platform and Distribution :colab Linux
• TensorFlow version:2.4.1
• TensorFlow-Addons version:0.12.1
• Python version:3.7.10
• Is GPU used? (yes/no): yes

I have used the specific versions of TensorFlow, TensorFlow-Addons, and Python as mentioned here : https://github.com/tensorflow/addons#python-op-compatibility-matrix

Describe the bug I am working on developing a neural machine translation system using the Seq2Seq model. I followed this tutorial from TensorFlow

https://www.tensorflow.org/addons/tutorials/networks_seq2seq_nmt

After training the model I want to save the model and use it elsewhere, Whoever I am facing the following issues with the usage of @tf.function.

#USING THE @tf.function in the decoder gives the error in training

TypeError: An op outside of the function building code is being passed a "Graph" tensor. It is possible to have Graph tensors leak out of the function building context by including a tf.init_scope in your function building code. For example, the following function will fail: @tf.function def has_init_scope(): my_constant = tf.constant(1.) with tf.init_scope(): added = my_constant * 2 The graph tensor has name: LuongAttention/memory_layer/Tensordot:0

#NOT USING THE @tf.function in the decoder trains the model but gives the error in saving

AssertionError: Tried to export a function which references untracked resource Tensor("LuongAttention/memory_layer/Tensordot:0", shape=(32, 11, 1024), dtype=float32).TensorFlow objects (e.g. tf.Variable) captured by functions must be tracked by assigning them to an attribute of a tracked object or assigned to an attribute of the main object directly.

Trackable Python objects referring to this tensor (from gc.get_referrers, limited to two hops): <tensorflow_addons.seq2seq.attention_wrapper.LuongAttention object at 0x7fad2f922590>

#NOT USING THE @tf.function in the decoder trains the model and running the model once, saves the model but gives the error after loading

TypeError: The argument 'cell' (<tensorflow.python.keras.saving.saved_model.load.AttentionWrapper object at 0x7facd7bb59d0>) is not an RNNCell: 'output_size' property is missing, 'state_size' property is missing, 'get_initial_state' method is required.

Code to reproduce the issue here is the complete code to replicate the error: google colab link

https://colab.research.google.com/drive/1ktxyOCHZe_QZ92Nk4lqujfhnvWc4YH0s?usp=sharing

Any pointers in this direction are highly appreciated, thank you!

[Kojo7]

I am having similar issues, please any help would be very great.