diff --git a/lenet5.py b/lenet5.py
index 8fd5d31..86dbd56 100644
--- a/lenet5.py
+++ b/lenet5.py
@@ -14,6 +14,7 @@ import keras.layers as layers
 from keras.preprocessing.image import ImageDataGenerator
 from keras.utils.np_utils import to_categorical
 from keras.callbacks import TensorBoard
+from keras.models import load_model
 
 EPOCHS = 10
 BATCH_SIZE = 128
@@ -46,12 +47,12 @@ def read_mnist(images_path: str, labels_path: str):
 
     return features, labels
 
-def display_image(position):
+def display_image(dataset, position):
     """
     Display image at position of the given dataset.
     """
-    image = train['features'][position].squeeze()
-    plt.title('Example %d. Label: %d' % (position, train['labels'][position]))
+    image = dataset['features'][position].squeeze()
+    plt.title('Example %d. Label: %d' % (position, dataset['labels'][position]))
     plt.imshow(image, cmap=plt.get_cmap('gray_r'))
     plt.show()
 
@@ -90,9 +91,9 @@ def main():
 
     # Step 4:
     # Display some images
-    display_image(0)
-    display_image(1)
-    display_image(2)
+    # display_image(train, 0)
+    # display_image(train, 1)
+    # display_image(train, 2)
 
     # Step 5:
     # Plot information about the training data
@@ -171,15 +172,32 @@ def main():
 
     tensorboard = TensorBoard(log_dir="logs/{}".format(time()))
 
+    """
     model.fit_generator(train_generator, steps_per_epoch=steps_per_epoch, epochs=EPOCHS,
                     validation_data=validation_generator, validation_steps=validation_steps,
                     shuffle=True, callbacks=[tensorboard])
 
+    model.save('lenet5.h5')
+    """
+
+    model = load_model('lenet5.h5')
+
     # Step 9:
     # Print results
     score = model.evaluate(test['features'], to_categorical(test['labels']))
     print('Test loss:', score[0])
     print('Test accuracy:', score[1])
 
+    # display_image(test, 0)
+    # display_image(test, 1)
+    # display_image(test, 2)
+
+    outputs = model.predict(test['features'])
+
+    for i, item in enumerate(outputs[0]):
+        print("%d: %.3f" % (i, item))
+
+    print(test['labels'][0])
+
 if __name__ == "__main__":
     main()