What version of PyTorch are you on (print(torch.__version__), and what is the output of nvidia-smi?
Can you also run the following and report back the log:
$ CUDA_LAUNCH_BLOCKING=1 gdb python
(gdb) r your_script.py
# when you get segfault
(gdb) bt