title: tf.keras.Model save: "AssertionError: Tried to export a function which references untracked object Tensor"
body: <p>I'm running this Tensorflow NMT tutorial: <a href="https://github.com/tensorflow/addons/blob/master/docs/tutorials/networks_seq2seq_nmt.ipynb" rel="nofollow noreferrer">https://github.com/tensorflow/addons/blob/master/docs/tutorials/networks_seq2seq_nmt.ipynb</a></p>
  <p>When I try to save the decoder:
  <code>decoder.save('decoder')</code>, I get:</p>
  <p><code>AssertionError: Tried to export a function which references untracked object Tensor(&quot;LuongAttention/memory_layer/Tensordot:0&quot;, shape=(1024, 23, 256), 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.</code></p>
  <p>I also tried registering the LuongAttention object like:</p>
  <pre><code>attention_mechanism = tfa.seq2seq.LuongAttention(units=units, memory=None, memory_sequence_length=BATCH_SIZE*[max_length_input])
  custom_objects = {&quot;LuongAttention&quot;: attention_mechanism}
  with tf.keras.utils.custom_object_scope(custom_objects):
    decoder = Decoder(vocab_tar_size, embedding_dim, units, BATCH_SIZE, attention_mechanism)```
  </code></pre>