One problem could be that the ouput variable is not freed after every pass of the model. It’s value and respective activations are still present when the second epoch starts.
You can do training and validation inside a function, so that there variables go out of scope after their call, and the memory is freed. Look at ‘Don’t hold to tensors’ in this Frequently Asked Questions — PyTorch 1.9.0 documentation 3
And yes, as @ptrblck mentioned, you should put validation under the torch.no_grad() context so it’s gradients are not computed. You don’t need them.