# Why does my custom categorical cross-entropy explode with Keras while not with TensorFlow?

I'm trying to train various Keras models on Pascal VOC 2012 dataset. The particularity of this dataset is that there is a particular class used to label "ambiguous" regions. These pixels are meant to be ignored when computing the loss and accuracy. The pixel value on the annotation images for this class is $$255$$.

For this reason, I need to implement a custom categorical cross-entropy to ignore the pixels labelled $$255$$:

def restricted_categorical_crossentropy(y_true, y_pred, class_labels):
def get_valid_probabilities_and_labels(annotation_batch_tensor,
logits_batch_tensor,
class_labels):
valid_batch_indices = tf.to_int32(valid_labels_indices)

valid_labels_batch_tensor = tf.gather_nd(params=annotation_batch_tensor, indices=valid_batch_indices)
valid_probabilities_batch_tensor = tf.gather_nd(params=logits_batch_tensor, indices=valid_batch_indices)

return valid_labels_batch_tensor, valid_probabilities_batch_tensor

valid_labels, valid_predictions = get_valid_probabilities_and_labels(
annotation_batch_tensor=y_true,
logits_batch_tensor=y_pred,
class_labels=class_labels)

cross_entropy = K.categorical_crossentropy(valid_labels, valid_predictions)

# The number of elements is different during each step due to mask out regions: normalize entropy.
return K.mean(cross_entropy)


When I compile my model with Keras' native K.categorical_crossentropy, I get values near $$2.5$$ and the model is able to slowly optimize that. When I use the custom function above, however, I get values like $$11563889.6988$$ and the model oscillates much more.

How does this make any sense granted

1. my custom function uses K.categorical_crossentropy on a subset of pixels of the same batches;
2. that exact same function works well when I'm not using Keras (I get values close to $$2.5$$ and lower)?

Also note that I preprocess the batches the exact same way in both situations (Keras vs. TensorFlow). Virtually nothing changed within the data pipeline except that I'm compiling a Keras model instead of running the session myself. I'm using the same TF records, same preprocessing function, same batching process.