我可以在RNNCell的__call__中存储状态吗

import tensorflow as tf

class MyRNNCell(tf.keras.layers.Layer):
    def __init__(self, hidden_dim):
        super(MyRNNCell, self).__init__()
        self.hidden_dim = hidden_dim
        self.state = None

    def build(self, input_shape):
        self.kernel = self.add_weight(shape=(input_shape[-1] + self.hidden_dim, self.hidden_dim),
                                      initializer='glorot_uniform',
                                      name='kernel')
        self.bias = self.add_weight(shape=(self.hidden_dim,),
                                    initializer='zeros',
                                    name='bias')

    def call(self, inputs):
        if self.state is None:
            self.state = tf.zeros([inputs.shape[0], self.hidden_dim])
        concat_inputs = tf.concat([inputs, self.state], axis=-1)
        output = tf.matmul(concat_inputs, self.kernel) + self.bias
        self.state = output  # 更新状态
        return output