Python Keras/Tensorflow自定义损失函数

我有一个用Keras构建的神经网络，我正在尝试训练它。输出层有4个节点。对于我试图解决的问题，我只想根据真值计算单个输出节点上的梯度。基本上，y_true看起来像[0,0,2,0]，其中零表示应该忽略的节点。然而，y_pred的形式为[1.2,3.2,4.5,6]。我希望在mse中只考虑第三个索引。这将要求我将y_pred中的索引0、1和3归零。我还没有找到一个合适的方法来做这件事

下面是我尝试过的代码，但它从loss函数返回NaN

def custom_mse(y_true, y_pred):

    return K.mean(K.square(tf.truediv(y_pred*y_true,y_true)-y_pred), axis=-1)

有没有办法对这些张量对象执行此简单操作？

如下操作：

[1.2,3.2,4.5,6]*[0,0,2,0] = [0,0,9,0]
[0,0,9,0]/2 = [0,0,4.5,0]

然后正常继续

这是执行此操作的代码：

def custom_mse(y_true, y_pred):
     return K.mean(K.square(tf.divide(tf.multiply(y_pred, y_true),tf.reduce_max(y_true))-y_true), axis=-1)

这样做：

[1.2,3.2,4.5,6]*[0,0,2,0] = [0,0,9,0]
[0,0,9,0]/2 = [0,0,4.5,0]

然后正常继续

这是执行此操作的代码：

def custom_mse(y_true, y_pred):
     return K.mean(K.square(tf.divide(tf.multiply(y_pred, y_true),tf.reduce_max(y_true))-y_true), axis=-1)