有没有办法在总体损失函数中增加基于keras‘自定义层’的/特定惩罚？

import tensorflow as tf
from tensorflow import keras

# 自定义层
class CustomLayer(keras.layers.Layer):
    def __init__(self, penalty_factor, **kwargs):
        super(CustomLayer, self).__init__(**kwargs)
        self.penalty_factor = penalty_factor

    def call(self, inputs):
        # 自定义层的前向传播逻辑
        return inputs

    def compute_output_shape(self, input_shape):
        return input_shape

# 自定义损失函数
def custom_loss(y_true, y_pred):
    # 自定义层的输出
    custom_output = model.get_layer('custom_layer').output

    # 计算自定义惩罚项
    penalty = tf.reduce_mean(tf.square(custom_output - y_true))

    # 总体损失函数
    total_loss = keras.losses.mean_squared_error(y_true, y_pred) + penalty_factor * penalty

    return total_loss

# 创建模型
input_layer = keras.layers.Input(shape=(input_shape,))
custom_layer = CustomLayer(penalty_factor)(input_layer)
output_layer = keras.layers.Dense(units=output_units)(custom_layer)
model = keras.models.Model(inputs=input_layer, outputs=output_layer)

# 编译模型
model.compile(optimizer='adam', loss=custom_loss)

# 训练模型
model.fit(x_train, y_train, epochs=epochs, batch_size=batch_size)