Solved by ptrblck in post #2 
              
              
                nn.CrossEntropyLoss expects raw logits as the model output, so remove the F.softmax in your forward method and check, if this would yield non-zero gradients.