Использование вывода модели массива numpy в определении графа тензорного потока

Question

У меня есть модель, которая ожидает вектор numpy на входе и возвращает вектор numpy на выходе. Сейчас я пытаюсь адаптировать некоторый существующий код, написанный для тензорных моделей, для работы с моей моделью. В какой-то момент, на этапе определения графика, у меня есть:

output = model.predict(new_point)

Есть ли способ обернуть предсказание модели (numpy array) внутри этой операции?

В настоящее время я пытаюсь с tf.py_function как:

@tf.function
def predict(self, data):
    tf.py_function(
        self.model.predict_proba,
        data,
        Tout=tf.float32
    )

Но я сталкиваюсь с: Expected list for 'input' argument to 'EagerPyFunc' Op, not Tensor