code: def masked_softmax_cross_entropy(preds, labels, mask):
      """Softmax cross-entropy loss with masking."""
      loss = tf.nn.softmax_cross_entropy_with_logits(
          logits=preds, labels=labels
      )
      mask = tf.cast(mask, dtype=tf.float32)
      mask /= tf.reduce_mean(mask)
      loss *= mask
      return tf.reduce_mean(loss)
  
label_desc: Use image node_id 1 for calling a global function with example usage: masked_softmax_cross_entropy(preds, labels, mask) and returns: tf