在@ tf.function模式下实例方法的第二次传递期间发生TypeError

Question

我写了一个定制的图层，它是Keras.layers.Conv2D的子类。

请参见this link

然后，我通过以下代码测试了该层：

c = MaskedConv2D(128,7,mask_type='A', padding='SAME')
test_data = tf.Variable(np.random.randn(1,28,28,3),dtype=tf.float32)
c(test_data)
test_data = tf.Variable(np.random.randn(1,28,28,3),dtype=tf.float32)
c(test_data)

输出为：

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: mul:0

我注释掉@tf.function时未发生错误。 在仍然使用@tf.funciton装饰器的同时如何解决此错误？