model_non_linear = MNIST_NN(input_shape=28*28,
 hidden_units=150,
 output_shape=len(class_names))