Description
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!