如何在Keras中实现自定义加权MSE损失函数？

from keras import backend as K
from keras.losses import mean_squared_error

def weighted_mse(y_true, y_pred):
    weights = K.constant([0.5, 0.5])  # 设置权重，这里假设有两个输出节点
    squared_difference = K.square(y_true - y_pred)
    weighted_squared_difference = K.dot(squared_difference, weights)
    return K.mean(weighted_squared_difference)

model.compile(optimizer='adam', loss=weighted_mse)