Thanks for the answer @James_Hickman!
I am also facing this problem right now, and am looking for a more efficient solution since I am dealing with large tensors.
Bumping this post in case anyone knows of a better solution now, because of the improvements in Pytorch over the last two years.