import tensorflow as tf

    else:
      return -tf.reduce_sum(log_sum_exp(log_probs), [1, 2])


def mse_loss(pred, labels):
  try:
    batch_size = tf.cast(pred.shape[0], tf.float32)
  except Exception as e:
    print('Pred is a tf tensor %s' % str(e.message))
    batch_size = tf.cast(tf.shape(pred)[0], tf.float32)
  loss_val = tf.sqrt(2 * tf.nn.l2_loss(pred - labels)) / batch_size
  return loss_val


def pullaway_loss(embeddings, name='pullaway_loss'):
  """Pull Away loss calculation.

  Args:
