import tensorflow as tf

parser.add_argument('--batch_size', type=int, default=32, help='batch size')
parser.add_argument('--max_train_step', type=int, default=50000, help='the maximum training step')
parser.add_argument('--model_path', type=str, default='', help='the path of checkpoint file')
args = parser.parse_args()

def model():
    x = tf.placeholder(tf.float32, [None, 784], name='x')
    gt = tf.placeholder(tf.float32, [None, 10], name='groundtruth')
    with tf.variable_scope('layer1'):
        w1 = tf.get_variable('weight1', [784, 1024], initializer=tf.random_normal_initializer())
        b1 = tf.get_variable('bias1', [1024], initializer=tf.constant_initializer(0.0))
        h1 = tf.nn.relu(tf.matmul(x, w1) + b1)
