Собственная функция потерь, tf.edit_distance, Ошибка в modle.compile

Question

Я пытаюсь реализовать свою собственную функцию потерь для сквозной сети обработки текстов (Keras + tenorflow). Целью было представить потерю с помощью WER (Word Error Rate). Первым шагом было определить расстояние редактирования и использовать его как потерю. Для этого я нашел tf.edit_distance после долгого поиска (у меня не было кода для саморазвития, чтобы запустить [обычный python ничего с keras или tf]). После некоторых попыток и ошибок и некоторых исследований я не получаю никаких сообщений об ошибках от своей функции потери напрямую, но теперь новое, с которым я ничего не могу сделать.

Я использую Keras 2.0.4 с tf 2.0a в Google Colab, во время работы Python 2 GPU.

Это определение моей функции потерь.

def create_sparse(dense):
  zero = K.tf.constant(0, dtype=K.tf.float32)
  where = K.tf.not_equal(dense, zero)
  indices = K.tf.where(where)
  values = K.tf.gather_nd(dense, indices)
  shape=K.tf.cast(K.tf.shape(dense), K.tf.int64)
  sparse = K.tf.SparseTensor(indices, values, shape)
  return sparse

def ed(y_true,y_pred):
  loss = K.tf.edit_distance(create_sparse(y_pred), create_sparse(y_true), normalize=True)
  # print(loss)
  return loss

def editDistanceLoss():
  def myHelperLoss(y_true, y_pred):
    return ed(y_true, y_pred)
  return myHelperLoss

И я использовал это так:

loss_editDistance = editDistanceLoss()
...
model.compile(loss=loss_editDistance, optimizer=optimizer, metrics=['accuracy'])

при запуске компиляции появляется следующее сообщение об ошибке:

TypeErrorTraceback (most recent call last)
<ipython-input-98-5787fc39817d> in <module>()
      2 # opt = Adam(lr=LR)  # keep calm and reduce learning rate
      3 # model.compile(loss='categorical_crossentropy', optimizer=optimizer, metrics=['accuracy'])
----> 4 model.compile(loss=loss_editDistance, optimizer=optimizer, metrics=['accuracy'])
      5 # model.compile(loss=loss_WER, optimizer=optimizer, metrics=['accuracy'])

1 frames
/usr/local/lib/python2.7/dist-packages/keras/engine/training_utils.pyc in weighted(y_true, y_pred, weights, mask)
    418             weight_ndim = K.ndim(weights)
    419             score_array = K.mean(score_array,
--> 420                                  axis=list(range(weight_ndim, ndim)))
    421             score_array *= weights
    422             score_array /= K.mean(K.cast(K.not_equal(weights, 0), K.floatx()))

TypeError: range() integer end argument expected, got NoneType.

если я добавлю print(loss) к функции ed, я получу следующий вывод Tensor("loss_21/activation_1_loss/edit_distance:0", dtype=float32).

Я не понимаю, почему я получаю это сообщение об ошибке. Может, кто-нибудь подскажет, что я делаю неправильно.