Ok, I solved it. The GPU was running out of memory, and probably it didn’t even have enough memory to print
RuntimeError: cuda runtime error (2) : out of memory at /opt/conda/conda-bld/pytorch_1518243271935/work/torch/lib/THC/generic/THCStorage.cu:58
After freeing up GPU’s memory, the code ran fine.