import tensorflow as tf

    train_step = tf.train.AdamOptimizer(
        learning_rate=config.learning_rate).minimize(cost)

    correct_prediction = tf.equal(tf.argmax(pred_Y, 1), tf.argmax(Y, 1))
    accuracy = tf.reduce_mean(tf.cast(correct_prediction, dtype=tf.float32))

    # --------------------------------------------
    # step4: Hooray, now train the neural network
    # --------------------------------------------
    # Note that log_device_placement can be turned ON but will cause console spam.

    # Initializing the variables
    init = tf.initialize_all_variables()
    # Add ops to save and restore all the variables.
    saver = tf.train.Saver()
    best_accuracy = 0.0

    # sess = tf.InteractiveSession(config=tf.ConfigProto(log_device_placement=False))
if (FLAG == 'train') : # If it is the training mode
    with tf.Session() as sess:
        # tf.initialize_all_variables().run()
        sess.run(init)  # .run()
        f.write("---Save model \n")

