This might be a good place to start: http://pytorch.org/docs/master/notes/faq.html#my-model-reports-cuda-runtime-error-2-out-of-memory 58
It seems that the first function keeps the history of ident for each for-loop iteration, which causes the CUDA memory error.