def METHOD_NAME(self, s):
    ss = ops.modal_dot(s, s, transpose_a=True)
    loss = -tf.linalg.trace(tf.math.sqrt(ss))
    if self.normalized_loss:
        n = float(tf.shape(s, out_type=tf.int32)[-2])
        c = float(tf.shape(s, out_type=tf.int32)[-1])
        loss = loss / tf.math.sqrt(n * c)
    return loss