loss.py

loss_objective_function = SparseCategoricalCrossentropy(from_logits=True, reduction='none')

def padded_loss_function(real, prediction):
    mask = tf.math.logical_not(tf.math.equal(real, 0))
    loss = loss_objective_function(real, prediction)

    mask = tf.cast(mask, dtype=loss.dtype)
    loss *= mask

    return tf.reduce_mean(loss)

training_loss = Mean(name='training_loss')
training_accuracy = SparseCategoricalAccuracy(name='training_accuracy')