tensorflow: TypeError: An op outside of the function building code is being passed a "Graph" tensor

System information

• Have I written custom code (as opposed to using a stock example script provided in TensorFlow): Yes
• OS Platform and Distribution: Windows 10
• TensorFlow installed from (source or binary): Source
• TensorFlow version (use command below): 2.4-RC0
• Python version: 3.8
• CUDA/cuDNN version: 11/8.0.4
• GPU model and memory: RTX 3080, 10GB

Describe the current behavior I am trying to build a custom RNN cell, based on this paper:

Using Fast Weights to Attend to the Recent Past Jimmy Ba, Geoffrey Hinton, Volodymyr Mnih, Joel Z. Leibo, Catalin Ionescu NIPS 2016, https://arxiv.org/abs/1610.06258

However, my implementation returns this error, but I have no clue what this means:

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: FW-RNN/while/fw_rnn_cell_2/add_2:0

Describe the expected behavior The model should simply start training.

Standalone code to reproduce the issue Code: https://gist.github.com/ion-elgreco/cc6fed29a4f6fb6813b71d5e7e8f2c87#file-fw_rnn-py

Other info / logs Error log: https://gist.github.com/ion-elgreco/cc6fed29a4f6fb6813b71d5e7e8f2c87#file-fw_rnn_error-txt