PyTorch, MLflow & Optuna: Experiment Tracking and Hyperparameter Optimization

The last two layers are fully-connected and a dropout is applied on each of them. Note that the dropout probability for the layers is passed into the network's constructor. This enables us to flexibly pass different values for this hyperparameter, as we will see in a bit. The final layer contains a logarithmic softmax function which gives us the probability for each of the 10 digit classes from MNIST. The digit with the highest probability is the one our network thinks is most likely visible on the input picture.