Could you wrap your testing code into a with torch.no_grad() block?
This would avoid storing the intermediate tensors, which would be needed to calculate the gradients in the backward pass.
Also, you might want to define a training and validation function, if that’s not already the case.
Since Python uses function scoping, some variables from the training might be freed additionally.