Oh sorry I did not see the zero grad. Glad I could help. Also another small thing. It doesn’t really matter but sometimes it looks better to define the loss function before the loop and just use it in the loop like this:
criterion = nn.CrossEntropyLoss()

outside the for loop and then in the for loop you would only have to do this
loss = criterion(output,labels)