I tried running your program with gdb, I’m getting a deadlock on the autograd execution engine. There are some posts related to this one, and a currently opened issue on PyTorch’s GitHub.
It seems to me that people have solved this by changing the environment, it could be related to PyTorch’s build and dependencies.