Refactored to be able to choose between two Models.
One is a simple fully connected multi-layer perceptron and one is a cnn