Hi KFrank,
You are correct. I guess that’s what happens when you’re new to PyTorch haha. I printed the wrong grads. So, I managed to have some plots for all gradients on the classifier, with tensorboard. Again, the model I used is VGG-16 with batch norm and already pre-trained weights. I freezed all the feature layers, and trained the layers on the classifier. And by seeing the values I’m getting suspicious.
Hyperparameters used for this one:

num_epochs = 10
batch_size = 64
lr = 0.01
SGD optimizer with default settings

I noticed that if I lower the batch size to 16 or less I get NaN values as output.
Here are the gradients histograms for the 3 layers of the classifier (0, 3 and 5, the layers in between are dropout and relu that’s why they’re not visualized):

image850×598 44.5 KB

And here are the gradients distributions for the 3 layers of the classifier:

image795×599 40.3 KB

This is pretty damn low, or I’m wrong? Seems like a vanishing gradient problem to me, but I’m not sure.