from tensorflow.python.keras.utils import losses_utilskl = tf.keras.losses.KLDivergence( reduction = losses_utils.ReductionV2.NONE, name = 'kullback_leibler_divergence') kl_loss = tf.reduce_mean(kl(logit1, logit2))