如何在Keras后台处理对角线元素？

import tensorflow as tf
from keras import backend as K

def process_diagonal(x):
    diagonal = tf.linalg.diag_part(x)  # 获取对角线元素
    processed_diagonal = tf.zeros_like(diagonal)  # 将对角线元素设置为零
    processed_x = tf.linalg.set_diag(x, processed_diagonal)  # 将处理后的对角线元素重新设置回张量
    return processed_x

from keras.layers import Input, Lambda

input_tensor = Input(shape=(10, 10))  # 假设输入张量的形状为(10, 10)
processed_tensor = Lambda(process_diagonal)(input_tensor)  # 应用自定义函数到输入张量上