Try to use the latest PyTorch (1.10).
The error indicates an out of bound memory access similar to a segfault on the CPU, something like an indexing error in the low level code.
You do reproduce this with the exact code you posted?
Best regards
Thomas