Use main.py
to visualize how the logit values that are passed into the softmax function change over time the model is trained with SGD (stochastic gradient descent) or the Adam optimizer.
This code if from my YouTube tutorial video called Softmax Function Explained In-Depth with 3D Visuals: