如何为Tensorflow LSTM类编写自定义调用函数？

import tensorflow as tf

class CustomLSTM(tf.keras.layers.Layer):
    def __init__(self, units):
        super(CustomLSTM, self).__init__()
        self.units = units

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

    def call(self, inputs, states):
        prev_output = states[0]
        h = tf.matmul(inputs, self.kernel)
        output = h + tf.matmul(prev_output, self.recurrent_kernel) + self.bias
        output = tf.nn.tanh(output)
        return output, [output]

# 使用自定义的LSTM层
lstm = CustomLSTM(units=64)
inputs = tf.keras.Input(shape=(10, 32))
x = tf.keras.layers.RNN(lstm)(inputs)
model = tf.keras.Model(inputs=inputs, outputs=x)