import tensorflow as tf

    # Memory allocation on the GPU as needed
    config = tf.ConfigProto()
    config.gpu_options.allow_growth = True
    config.allow_soft_placement = True
    
    # Train
    with tf.Session(config=config) as sess:
        try:
            summary_writer = tf.summary.FileWriter(tensorboard_dir, sess.graph)
            
            sess.run(tf.global_variables_initializer())
            
            # saved model restoring
            if args.restore:
                # Restore saved model if the user requested it, default = True
                try:
