входные параметры tf.function

Question

Я пишу функцию в tenorflow 2 и использую tf.keras для написания модели. Функция определена ниже:

@tf.function
def mask_output(input_tensor,mask):
    if tf.math.count_nonzero(mask) > 0:
        output_tensor = tf.multiply(input_tensor, mask)
    else:
        output_tensor = input_tensor
    return output_tensor

Два параметра, которые я даю, это тензор в модели. Однако, когда я определяю модель и вызываю эту функцию в определении модели, он говорит:

{_ SymbolicException} Входные данные для активной функции выполнения не могут быть символами Keras символами c тензорами, но найдены

[<tf.Tensor 'a_dense2/Identity:0' shape=(None, 12, 5) dtype=float32>, <tf.Tensor 'a_mask_input:0' shape=(None, 12, 5) dtype=float32>]

Как это решить? Почему я не могу назвать эту функцию с тензорным входом двух кератов?