Custom Neural Network Implementation on MNIST using Tensorflow 2.0?

Question

I tried to write a custom implementation of basic neural network with two hidden layers on MNIST dataset using *TensorFlow 2.0 beta* but I\'m not sure what went wrong here but m

Answer 1

Also If there's something I could improve in the code do let me know as well.

Embrace the high-level API for something like this. You can do it in just a few lines of code and it's much easier to debug, read and reason about:

(x_train, y_train), (x_test, y_test) = tfds.load('mnist', split=['train', 'test'], 
                                                  batch_size=-1, as_supervised=True)

x_train = tf.cast(tf.reshape(x_train, shape=(x_train.shape[0], 784)), tf.float32)
x_test  = tf.cast(tf.reshape(x_test, shape=(x_test.shape[0], 784)), tf.float32)

model = tf.keras.models.Sequential([
  tf.keras.layers.Dense(512, activation='sigmoid'),
  tf.keras.layers.Dense(256, activation='sigmoid'),
  tf.keras.layers.Dense(10, activation='softmax')
])
model.fit(x_train, y_train, epochs=5)
model.evaluate(x_test, y_test)