Having memory overflow issues

Hi i have been trying to build a simple CNN for the MNIST dataset, but about after 4 epochs, I have a memory overflow issues. Any suggestions on how I can fix this? Thank you!

notebook here : https://colab.research.google.com/drive/16RFi_fERIyHMpiNXlGJb01faUsJWrkot

Hi Nishanth,
The notebook doesn’t have the code for the training loop. Could you upload that as well? Since, the training loop is the one causing the error.

sorry, I accidentally removed it. It should be there now thought!

thanks for any help

Hi Nishanth,
I had a look at the training code. The training loop is correct. The validation had this line:

loss += criterion(outputs, classes)

This was causing pytorch to retain the memory of all the previous graphs executed for validation.
Changing it to the below line will solve your problem.

loss += criterion(outputs, classes).item()

Yep! that solved my problem. Thank you!

1 Like