Tensorflow custom cost function: NaN/Inf values despite clipping_by_value

I'm trying to play around with a custom cost function on the MNIST data set (28x28 images of handwritten digits from 0 to 9). I use this simple architecture:

model = tf.keras.models.Sequential([
  tf.keras.layers.Flatten(input_shape=(28, 28)),
  tf.keras.layers.Dense(128, activation='relu'),
  tf.keras.layers.BatchNormalization(),
  tf.keras.layers.Dense(10, activation='softmax')
])

and eventually fit my model with my custom cost function

model.compile(optimizer = 'adam',
    loss = custom),
    metrics = ['accuracy'])

My custom cost function is fairly simple (code below) but it results in unexpected NaN's...

def custom(target, pred):
    # Slight variation on usual crossentropy loss: in addition to
    # the term -log(y_) for y_ the probability corresponding to the
    # true label, we add -log(1 - A) with A the L^10 norm of the
    # vector of probabilities corresponding to the incorrect labels.
    # 10 is arbitrary, for the purpose of this post. The higher the
    # value, the faster I get NaN's.
    #
    # target: labels, uint8 between 0 and 9, shape (, 1)
    # pred: predictions, probabilities between 0 and 1, shape (, 10)

    k = 10 # num classes

    # turning vector of labels into one-hot matrix
    target_oh = tf.one_hot(tf.squeeze(target), k) # same shape as pred

    # applying the target_oh mask.
    correct = tf.math.multiply(target_oh, pred) # predictions on the correct labels
    wrong = tf.math.subtract(pred, correct) # predictions on the wrong labels

    # first clipping (should be useless?)
    _epsilon = tf.keras.backend.epsilon()
    correct = tf.clip_by_value(correct, _epsilon, 1. - _epsilon)
    wrong = tf.clip_by_value(wrong, _epsilon, 1. - _epsilon)
    
    # compute score for correct and wrong labels (we hope to be close to 1).
    correct_score = tf.norm(correct, ord=10, axis=1)
    wrong_score = tf.subtract(tf.constant(1.), tf.norm(wrong, ord=10, axis=1)) # 1 - (norm of wrong labels' probabilities)

    # second clipping: in bad case scenario, one of the scores is close to 0 and log explodes
    correct_score = tf.clip_by_value(correct_score, _epsilon, 1. - _epsilon)
    wrong_score = tf.clip_by_value(wrong_score, _epsilon, 1. - _epsilon)

    # now taking log should be fine
    correct_cost = tf.math.negative(tf.math.log(correct_score))
    wrong_cost = tf.math.negative(tf.math.log(wrong_score))

    # adding costs
    cost = tf.math.add(correct_cost, wrong_cost)

    return cost

After only a few steps of computing, I get

Epoch 2/10 1/1 [==============================] - 0s 266ms/step - loss: nan - accuracy: 0.0987

Upon trying to use tf.debugging.enable_check_numerics(), I get that +Inf is passed to Pow at some point, but I really don't understand how it is possible with my clipping in custom:

: Tensor had +Inf values [[{{node gradient_tape/lambda/norm/Pow_1/Pow/CheckNumericsV2}}]] [Op:__inference_train_function_1944]

If anyone can help me, I'd be very grateful <3 Thanks for reading!!