Well, you obviously have a much better GPU than mine 
In any case, even after updating the GPU driver and CUDNN version, the program still gets stuck in the second iteration. Is there any kind of debugging that I can perform that will allow a better understanding of the problem?
Will be glad to receive your guidance.
EDIT: it seems like the program now gets stuck before the “optimizer.step()” in the second iteration. The reason for the error is: “denom = exp_avg_sq.sqrt().add_(group[‘eps’])” in the Adam optimizer routine.
Thanks a lot again.