relational_memory.py: InvalidArgumentError: node has inputs from different frames.

[songpipi]

I train a GAN model with relational rnn as my discriminator, I use tensorflow 1.12.0 and python 2.7. I use tf.contrib.framework.gan.gan_model() as my GAN model. Part of code :

from sonnet.python.modules import relational_memory
...
def discriminator():
  ...
  cell = relational_memory.RelationalMemory( mem_slots = 1, head_size = 256, num_heads = 2)
  rnn_out, state = tf.nn.dynamic_rnn(cell, rnn_inputs, length, dtype=tf.float32)
  ...

The code works, it can save model.ckpt-0, but when the data starts coming in, it raises: InvalidArgumentError: node train/discriminator_train/gradients/Discriminator_1/rnn/while/Switch_3_grad/b_switch (defined at /data/songpp/anaconda2/envs/tf12py2/lib/python2.7/site-packages/tensorflow/contrib/training/python/training/training.py:447) has inputs from different frames. The input node train/discriminator_train/gradients/Discriminator_1/rnn/while/Switch_3_grad_1/NextIteration (defined at /data/songpp/anaconda2/envs/tf12py2/lib/python2.7/site-packages/tensorflow/contrib/training/python/training/training.py:447) is in frame ''. The input node train/discriminator_train/gradients/Discriminator_1/rnn/while/Exit_3_grad/b_exit (defined at /data/songpp/anaconda2/envs/tf12py2/lib/python2.7/site-packages/tensorflow/contrib/training/python/training/training.py:447) is in frame 'train/discriminator_train/gradients/Discriminator_1/rnn/while/while_context'.

I tried to replace the discriminator with LSTM, and there was no problem. Why the problem and What should I do? Please help me out, thank you!