tenorflow 2.2 простая пользовательская потеря не работает (но работает в tf2.1)

Question

У меня есть следующая простая реализация настраиваемой взвешенной потери кроссэнтропии в тензорном потоке 2 (извините за неоптимальное форматирование, не смог заставить его работать должным образом):

class Weighted_Crossentropy_Loss(tf.keras.losses.Loss):

def __init__(self, class_weights, reduction=tf.keras.losses.Reduction.NONE, name=None, label_smoothing=0):
    super(Weighted_Crossentropy_Loss).__init__()
    self.n_class = len(class_weights)
    self.class_weights = tf.constant(np.array(class_weights, dtype=np.float32))
    self.reduction = reduction
    self.name = name
    self.label_smoothing = label_smoothing
    self.loss = tf.keras.losses.BinaryCrossentropy(from_logits=True, label_smoothing=label_smoothing,
                                                   reduction=reduction)

# @tf.function
def call(self, labels, logits):
    # adapted from Christian Baumgartner's code from the acdc challenge
    flat_logits = tf.reshape(logits, [-1, self.n_class])
    flat_labels = tf.reshape(labels, [-1, self.n_class])
    weight_map = tf.multiply(flat_labels, self.class_weights)
    weight_map = tf.reduce_sum(weight_map, axis=-1)

    loss_map = self.loss(y_true=flat_labels, y_pred=flat_logits)
    weighted_loss = tf.multiply(loss_map, weight_map)
    return weighted_loss`

Использование указанной выше потери для обучения модели в пользовательском цикле обучения с использованием tf2.1 работает нормально, но в tf2.2 я получаю сообщение об ошибке

AttributeError: 'Weighted_Crossentropy_Loss' object has no attribute '_name_scope'

при фактическом вызове функции потерь.

Затем просто для этого добавив self._name_scope = 'loss'в приведенную выше функцию инициализации, чтобы увидеть, что происходит, ошибка становится

AttributeError: 'Weighted_Crossentropy_Loss' object has no attribute '_allow_sum_over_batch_size'

Произошли ли большие изменения в том, как определять класс потерь в тензорном потоке? Обратите внимание, что я не использую какую-либо стратегию распространения или что-то подобное.