高级API中的Tesorflow自定义层: throws对象没有属性“”_expects_mask_arg“”错误

import tensorflow as tf

class CustomLayer(tf.keras.layers.Layer):
    def __init__(self, units, expects_mask_arg=False, **kwargs):
        self.units = units
        self.expects_mask_arg = expects_mask_arg
        super(CustomLayer, self).__init__(**kwargs)

    def build(self, input_shape):
        # 定义层的参数
        self.kernel = self.add_weight("kernel", shape=(input_shape[1], self.units), initializer="random_normal", trainable=True)
        super(CustomLayer, self).build(input_shape)

    def call(self, inputs, mask=None):
        # 定义层的计算逻辑
        output = tf.matmul(inputs, self.kernel)
        return output

    def compute_mask(self, inputs, mask=None):
        # 定义掩码计算逻辑
        if self.expects_mask_arg:
            return mask
        else:
            return None

# 使用自定义层
input_tensor = tf.keras.Input(shape=(10,))
custom_layer = CustomLayer(20, expects_mask_arg=True)
output_tensor = custom_layer(input_tensor)

input_tensor = tf.keras.Input(shape=(10,))
custom_layer = CustomLayer(20, expects_mask_arg=True)
output_tensor = custom_layer(input_tensor)