Tensorflow аннотирует универсальный кодировщик предложений с помощью функции tf.function

Question

У меня есть некоторый рабочий код, который загружает Universal Sentence Encoder, который был загружен из TF HUB, я загружаю его таким образом, чтобы я мог использовать TF2.

model = tf.saved_model.load('./use_lite_saved_model')

Затем мы создаем * Модель 1004 * также необходима, так как я использую облегченную версию:

sp = spm.SentencePieceProcessor()
sp.Load('./use_lite_saved_model/assets/universal_encoder_8k_spm.model')

Затем у нас есть служебная функция, которая использует sentencepiece для создания ввода, необходимого для USE Lite, я оставляю его для краткости, использование:

values, indices, dense_shape = process_to_IDs_in_sparse_format(sp, df['page_doc'])

Затем мы используем модель:

tensors_out_dict = model.signatures['default'](
    values=values,
    indices=indices,
    dense_shape=dense_shape
)
tensor_out = tensors_out_dict['default']

Все это прекрасно работает! Проблема в том, что, пытаясь улучшить производительность при выполнении описанной выше операции для многих page_docs, я пытался использовать аннотацию tf.function для использования AutoGraph, но без удачи каждая попытка выдает ошибку, окончательная версия:

@tf.function
def embed(values, indices, dense_shape):
    tensors_out_dict = model.signatures['default'](
        values=values,
        indices=indices,
        dense_shape=dense_shape
    )
    tensor_out = tensors_out_dict['default']
    return tensor_out

Здесь я аннотировал только часть кода предсказания, текущая ошибка:

FailedPreconditionError:  Attempting to use uninitialized value StatefulPartitionedCall/Encoder_en/KonaTransformer/Encode/Layer_1/TransformerLayer/MultiheadAttention/output_transform_single/kernel/part_0
     [[{{node StatefulPartitionedCall/Encoder_en/KonaTransformer/Encode/Layer_1/TransformerLayer/MultiheadAttention/output_transform_single/kernel/part_0/read}}]] [Op:__inference_embed_31861]

Function call stack:
embed