Y_true = np.argmax(y_test,axis = 1)
confusion_mtx = confusion_matrix(Y_true, prediction_values)
sns.heatmap(confusion_mtx, annot=True, fmt="d")
plt.ylabel('True')
plt.xlabel('Predicted')
# save weights to HDF5
model.save_weights("CNN_tf.keras_mnist.h5")