import tensorflow as tf

            if not config.aug_same and config.aug:
                recon_feat = tf.concat([features, target['aug']], -1)
                print('Use recon feature ', name, recon_feat)
                logprob = heads[name](recon_feat).log_prob(target[name])
                # logprob = heads[name](features).log_prob(target['ori_img'])
            else:
                logprob = heads[name](features).log_prob(target[name])
            objectives.append(Objective(name, logprob, max, include, exclude))

    objectives = [o._replace(value=tf.reduce_mean(o.value)) for o in objectives]

    return objectives, cstr_pct



def contra_step_lossV1(pred, tgt, temp=10.0):
    # Step-wise contrastive loss
    pred1, pred2 = tf.split(pred, 2, axis=0)
