Add code for later PyTorch versions

We could put some version parsing and appropriate if-guards into the
code, but I'm afraid of the increase in complexity.