Pytorch model training loss does not improve. Are the logistic regression model parameters/weights not updating?

Question

I am attempting to use the following model to classify images. The training loss does not appear to converge/improve. Can you check the code and see if this may be a model issue