Пользовательская потеря в Keras - медленная компиляция и соответствие

Question

Я пытаюсь выполнить функцию потери в Keras, как в Pytorch https://pytorch.org/docs/master/generated/torch.nn.MultiLabelMarginLoss.html

но на сборку model.compile уходит много времени, а после этого очень долго нужно ее обучать (особенно первую партию) (я использую multi gpu)

возможно это из-за циклов или неправильного использования Keras.backend,

Вот код: (L - количество классов)

def mlm_loss(y_true, y_pred):
loss=float(0)
a = tf.keras.backend.constant(1, dtype='float32')
for s in range(batch_size): # for each sample in batch
    for i in range(L):
        for j in range(L):
            loss=loss + y_true[s][i]*(a-y_true[s][j])*(a-(y_pred[s][i]-y_pred[s][j])) #two conditions
l= tf.keras.backend.constant(L, dtype='float32')            
loss=a/l*loss           
return loss

Спасибо за помощь