Как я могу обновить вес модели с двумя разными функциями потерь одновременно на TF2?

Я хочу обучить модель классификации с двумя потерями следующим образом:

model.compile(optimizer=adam)

@tf.function
def train(model, inputs_data_1, inputs_data_2, y):
  with tf.GradientTape(persistent=True) as tape:
    logits1, features1 = model(inputs_data_1)  # logits: output of fully-connected layer
    logits2, features2 = model(inputs_data_2)  # features: output of feature extractor

    loss_fn1 = cross-entropy(y, logits1)
    loss_fn2 = euclidean_dist(features1-features2)

    losses = loss_fn1 + loss_fn2

  optim.apply_gradients(zip(tape.gradient(losses, model.trainable_weights), model.trainable_variables))

когда я пытаюсь это сделать, он просто остановился без ошибки.

Я не менял входные данные с помощью tf.split или tf.reshape.

как я могу скомпилировать модель и тренироваться с двумя потерями?

Пожалуйста, дайте мне несколько мнений или реализацию кода, чтобы сослаться на эту проблему. Спасибо.