My model is working now. I believe due to large convolution filters (14 * 14) and channel sizes (512 and 10000) in my architecture, parameter count was too high and the GPU memory was blowing up. I still don’t understand why memory blew up exactly after loss.backward() method but after inspecting tensor sizes and model parameter count, GPU memory utilization seems reasonable.
Lastly I’m closing this discussion since solution is simply to reduce filter size to keep parameter count at conv3 low. Thanks for all your help.