Custom loss function with Keras - tensorflow

Question

I am trying to write a custom Keras loss function but I am having issues with implementing and debugging my code. My target vector is:
y_pred = [p_conf, p_class_1, p_class_2]
where, p_conf = confidence an event of interest was detected
y_true examples:
[0, 0, 0] = no event of interest
[1, 1, 0] = first class event
[1, 0, 1] = second class event
I get relatively good results using multi-label classification (i.e. using a sigmoid activation in my final layer and binary_crossentropy loss function) but I want to experiment and improve my results using a custom loss function that calculates the:
binary_crossentropy loss for when y_true = [0, ..., ...]
categorical_crossentropy loss for when y_true = [1, ..., ...]
This is a simplified loss function used by the YOLO object detection algorithm. I tried adapting an existing Keras / TensorFlow implementation of the YOLO loss function but have not been successful.
Here is my current working code. It runs but generates unstable results. i.e. loss and accuracy decreases over time. Any assistance would be greatly appreciated.
import tensorflow as tf
from keras import losses
def custom_loss(y_true, y_pred):
# Initialisation
mask_shape = tf.shape(y_true)[:0]
conf_mask = tf.zeros(mask_shape)
class_mask = tf.zeros(mask_shape)
# Labels
true_conf = y_true[..., 0]
true_class = tf.argmax(y_true[..., 1:], -1)
# Predictions
pred_conf = tf.sigmoid(y_pred[..., 0])
pred_class = y_pred[..., 1:]
# Masks for selecting rows based on confidence = {0, 1}
conf_mask = conf_mask + (1 - y_true[..., 0])
class_mask = y_true[..., 0]
nb_class = tf.reduce_sum(tf.to_float(class_mask > 0.0))
# Calculate loss
loss_conf = losses.binary_crossentropy(true_conf, pred_conf)
loss_class = tf.nn.sparse_softmax_cross_entropy_with_logits(labels=true_class, logits=pred_class)
loss_class = tf.reduce_sum(loss_class * class_mask) / nb_class
loss = loss_conf + loss_class
return loss