Обучите предварительно обученную пользовательскую модель с другой функцией потерь

В TF2 keras я обучил автоэнкодер, используя tensorflow.keras.losses.MeanSquaredError в качестве функции потерь. Теперь я хочу дополнительно обучить эту модель, используя другую функцию потерь, а именно tenorflow.keras.losses.KLDivergence. Причина этого в том, что изначально обучение без учителя проводится для репрезентативного обучения. Затем, имея сгенерированные вложения, я могу кластеризовать их и использовать эти кластеры для самоконтроля, то есть меток, позволяя вторую контролируемую потерю и улучшая модель дальше.

Это не переносное обучение как таковое, поскольку в модель не добавляются новые слои, изменяется только функция потерь, и модель продолжает обучение.

Я пробовал использовать предварительно обученную модель с потерей MSE в качестве свойства новой модели:

class ClusterBooster(tf.keras.Model):

 def __init__(self, base_model, centers):
    super(ClusterBooster, self).__init__()

    self.pretrained = base_model
    self.centers = centers


 def train_step(self, data):
    with tf.GradientTape() as tape:
        loss = self.compiled_loss(self.P, self.Q, regularization_losses=self.losses)

    # Compute gradients
    gradients = tape.gradient(loss, self.trainable_variables)
    
    # Update weights
    self.optimizer.apply_gradients(zip(gradients, self.trainable_variables))

    return {m.name: m.result() for m in self.metrics}

где потеря - это потеря KL между распределениями P и Q. Распределения вычисляются в функции обратного вызова вместо модели train_step, поскольку мне нужен доступ к текущей эпохе (P обновляется каждые 5 эпох, а не каждую эпоху):

def on_epoch_begin(self, epoch, logs=None):
    z = self.model.pretrained.embed(self.feature, training=True)
    z = tf.reshape(z, [tf.shape(z)[0], 1, tf.shape(z)[1]])  # reshape for broadcasting
    # CALCULATE Q FOR EVERY EPOCH
    partial = tf.math.pow(tf.norm(z - self.model.centers, axis=2, ord='euclidean'), 2)
    nominator = 1 / (1 + partial)
    denominator = tf.math.reduce_sum(1 / (1 + partial))
    self.model.Q = nominator / denominator

    # CALCULATE P EVERY 5 EPOCHS TO AVOID INSTABILITY
    if epoch % 5 == 0:
        partial = tf.math.pow(self.model.Q, 2) / tf.math.reduce_sum(self.model.Q, axis=1, keepdims=True)
        nominator = partial
        denominator = tf.math.reduce_sum(partial, axis=0)
        self.model.P = nominator / denominator

Однако при выполнении apply_gradients () я получаю:

ValueError: No gradients provided for any variable: ['dense/kernel:0', 'dense/bias:0', 'dense_1/kernel:0', 'dense_1/bias:0', 'dense_2/kernel:0', 'dense_2/bias:0', 'dense_3/kernel:0', 'dense_3/bias:0']

Я думаю, что это связано с тем, что предварительно обученная модель не настроена на дальнейшее обучение где-то внутри новой модели (вызывается только метод embed (), который не обучает модель). Это правильный подход, и я просто что-то упускаю, или есть способ лучше?