import tensorflow as tf

    self.d_loss += self._lambda * penalty_loss

    z_projs = tf.concat([z_projs_real, z_projs_fake], 0)
    z_aug_projs = tf.concat([z_aug_projs_real, z_aug_projs_fake], 0)

    sims_logits = tf.matmul(z_projs, z_aug_projs, transpose_b=True)
    logits_max = tf.reduce_max(sims_logits,1)
    sims_logits = sims_logits - tf.reshape(logits_max, [-1, 1])
    sims_probs = tf.nn.softmax(sims_logits)

    sim_labels = tf.constant(np.arange(bs * 2, dtype=np.int32))
    sims_onehot = tf.one_hot(sim_labels, bs * 2)
