Есть ли способ применить градиенты для нескольких выходных слоев в tf.GradientTape?

Question

Я пытаюсь применить градиенты к двум выходным моделям, но результат показывает, что модель не учится и потери не уменьшаются, мне нужно ваше поддерживаемое Спасибо.

@tf.function deftrain_step (inp, targ, intent, enc_hidden):

loss = 0
intent_loss = 0

with tf.GradientTape(persistent= True) as tape:

    enc_output, enc_hidden = encoder(inp, enc_hidden)

    dec_hidden = enc_hidden




    dec_input = tf.expand_dims([targ_lang.word_index['<start>']] * BATCH_SIZE, 1)

    # Teacher forcing - feeding the target as the next input
    for t in range(1, targ.shape[1]):

        # passing enc_output to the decoder
        predictions, dec_hidden, _ =slot_decoder(dec_input, dec_hidden, enc_output)
        intent_pred, _ = intent_decoder(dec_hidden, enc_output)

        loss += loss_function(targ[:, t], predictions)
        intent_loss = loss_function(intent, intent_pred)

        # using teacher forcing
        dec_input = tf.expand_dims(targ[:, t], 1)

batch_loss = (loss / int(targ.shape[1])) + intent_loss

intent_variables = encoder.trainable_variables + intent_decoder.trainable_variables
slot_variables = encoder.trainable_variables + slot_decoder.trainable_variables

intent_gradients = tape.gradient(intent_loss, intent_variables)
slot_gradients = tape.gradient(loss, slot_variables)


optimizer.apply_gradients(zip(intent_gradients, intent_variables))
optimizer.apply_gradients(zip(slot_gradients, slot_variables))

del tape
return batch_loss + intent_loss