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()