jabhinav:

I want to debug what is consuming most of the memory because the same error appears on GPU with larger memory of 32GBs (tried on k80 and v100).


I would recommend to add print(torch.cuda.memory_allocated()) statements in your code and check each operation or layer sequentially.
This would give you the overview where most of the memory is used.