import tensorflow as tf

        test = ContextAEReach()
    elif args.experiment_type == "push":
        test = ContextAEPush()
    elif args.experiment_type == "pushreal":    
        test = ContextAEPushReal()
    elif args.experiment_type == "sweep":
        test = ContextAESweep()
                
    test.build(tfinput, args.ablation_type)


    config = tf.ConfigProto()
    config.gpu_options.allow_growth=True
    sess = tf.Session(config=config)
    learning_rate = tf.placeholder(tf.float32, shape=[])
    optimizer = tf.train.AdamOptimizer(learning_rate).minimize(test.loss)
    sess.run(tf.global_variables_initializer())
    allloss = []
    validloss = []
    itr = 0
    saver = tf.train.Saver()

    n = vdata.shape[1]
    nlen = vdata.shape[0]
    ntrain = int(0.8*n)
    nvalid = n - ntrain
    validdata = vdata[:, ntrain:]
    traindata = vdata[:, :ntrain]
    while True:
        choicesrc = np.random.choice(ntrain, batch_size)
