Tensorflow with Hidden Layers

import tensorflow as tf
from tensorflow.examples.tutorials.mnist import input_data

mnist = input_data.read_data_sets(‘data/‘, one_hot=True)

num_classes = 10
batch_size = 64
hidden_units = 50
input_size = 784
train_iter = 10000

X = tf.placeholder(tf.float32, shape=[None, input_size])
y = tf.placeholder(tf.float32, shape=[None, num_classes])

W1 = tf.Variable(tf.random_normal(shape=[input_size, hidden_units], stddev=0.1))
b1 = tf.Variable(tf.constant(0.1, shape=[hidden_units]))
W2 = tf.Variable(tf.random_normal(shape=[hidden_units, num_classes], stddev=0.1))
b2 = tf.Variable(tf.constant(0.1, shape=[num_classes]))

a1 = tf.nn.relu(tf.matmul(X, W1) + b1)
a2 = tf.nn.relu(tf.matmul(a1, W2) + b2)

loss = tf.reduce_mean(tf.nn.softmax_cross_entropy_with_logits(logits=a2, labels=y))
train_op = tf.train.GradientDescentOptimizer(0.01).minimize(loss)

with tf.Session() as sess:
    sess.run(tf.global_variables_initializer())
    for step in range(10000):
        mini_batch = mnist.train.next_batch(batch_size)
        X_temp = mini_batch[0]
        y_temp = mini_batch[1]
        sess.run(train_op, feed_dict={X:X_temp, y:y_temp})
        if step % 1000 == 0:
            loss_var = sess.run(loss, feed_dict={X: X_temp, y: y_temp})
            print("loss:", loss_var)