如何使用functional API调用Keras中的自定义图层

from tensorflow.keras.layers import Layer
from tensorflow.keras import Input, Model

class CustomLayer(Layer):
    def __init__(self, units, activation=None):
        super(CustomLayer, self).__init__()
        self.units = units
        self.activation = activation

    def build(self, input_shape):
        self.kernel = self.add_weight("kernel", shape=(input_shape[1], self.units))
        self.bias = self.add_weight("bias", shape=(self.units,))

    def call(self, inputs):
        output = inputs @ self.kernel + self.bias
        if self.activation is not None:
            output = self.activation(output)
        return output

input_layer = Input(shape=(input_shape,))

custom_layer = CustomLayer(units=64, activation="relu")

output_layer = custom_layer(input_layer)

model = Model(inputs=input_layer, outputs=output_layer)

model.compile(optimizer="adam", loss="mse")
model.fit(x_train, y_train, epochs=10, batch_size=32)