Индекс трехмерного тензора с двумерным тензором в лямбда-слое в Керасе

Question

Я пытаюсь использовать 2-D тензор для индексации 3-D тензора в Tensorflow в слое Keras Lambda. Я использую следующую функцию для лямбда-слоя:

def testFunc(x):
   indexMatrix = K.cast(x[0], dtype = "int32")
   valueMatrix = x[1]
   idndex = tf.meshgrid(*[tf.range(s) for s in indexMatrix.get_shape()],indexing='ij')
   idx = tf.stack(index + [indexMatrix], axis=-1)
   output = tf.gather_nd(valueMatrix, idx)    
   return  output

и использую лямбда-слой, например:

layer = Lambda(testFunc, output_shape)([model.layers[0].output,model.layers[-1].output])

Когда я запускаю код, я получаю ошибку во время построения графика, потому что измерениеindexMatrix и valueMatrix неизвестны, а tf.range () нужны известные измерения. Ошибка

"Cannot convert an unknown Dimension to a Tensor: ?" 

Кто-нибудь знает, как я могу решить эту проблему? спасибо