Пользовательские потери в Tensorflow 2.0.0

Question

Я хочу реализовать пользовательские потери, которые можно рассчитать, используя каждый образец. Расчет потерь немного сложен и требует от меня использования внешнего файла Python для этого (или можно предположить, что мы даем входные данные для функции).

Как я могу это реализовать? Можно ли использовать @tf.function аннотацию и сделать ее графиком?

Вот как это должно выглядеть

def loss(input,output):
    loss = 0    
    for x, y in zip(input, output):
        sim = Class(x) 
        a = sim.GetA()
        b = sim.GetB()

        loss = loss + np.linalg.norm(np.dot(a,b)+y)
    return loss