norhther:

with torch.no_grad():


As a side note, with torch.no_grad():  should be for validation only. While training, you should not have the with torch.no_grad(): block. This may result in problems you defined as strange results.