def compute_bce_loss(logits, entity_bool, mask):
        """Compute binary cross-entropy loss for logits."""
        batch_labels = tf.expand_dims(entity_bool, 1)
        position_labels = batch_labels * tf.ones(
            (1, seq_length - 1), dtype=tf.float32)
        # Ignore [CLS] for this loss.
        loss_ = tf.nn.sigmoid_cross_entropy_with_logits(
            labels=position_labels, logits=logits[:, 1:])
        mask_f = tf.cast(mask[:, 1:], tf.float32)
        loss = tf.reduce_sum(loss_ * mask_f) / tf.maximum(
            1e-12, tf.reduce_sum(mask_f))
        return loss
Compute binary cross-entropy loss for logits.
language/labs/drkit/run_dualencoder_lsf.py