Дизайн Custom Keras Функция потери MSE

Question

Я хочу создать пользовательскую функцию потери MSE, которая использует только один элемент из выходного слоя.

То, что у меня сейчас есть, это:

def new_loss(y_true, y_pred):
   index_pred = K.argmax(K.abs(y_pred), axis = -1)
   pred = tf.gather(y_true, index_pred, axis = 1)

   index_true = K.argmin(K.abs(y_true), axis = -1)
   true = tf.gather(y_true, index_true, axis = 1)

   return K.mean(K.sqrt(K.square(K.log(pred)-K.log(true))))

но выдает ошибку,

An operation has `None` for gradient.

Я искал и не могу найти что-нибудь для работы в моем сценарии.

Я работаю с Keras и Tensorflow в качестве бэк-энда.

Заранее спасибо.

EDIT:

Я пытался

def new_loss(y_true, y_pred):
   index_pred = K.argmax(K.abs(y_pred), axis = -1)
   pred = tf.gather(y_pred, index_pred, axis = 1)

   index_true = K.argmin(K.abs(y_true), axis = -1)
   true = tf.gather(y_true, index_true, axis = 1)

   return K.mean(K.sqrt(K.square(K.log(pred)-K.log(true))))

И это не дает ошибки. Так что проблема не в K.argmax / K.argmin.