i think this was a bug that is now fixed in master and will be in the next release a week to two weeks away.
If you want an immediate fix, do consider building from master: https://github.com/pytorch/pytorch#from-source 8
We really apologize for having to make you build from source  for the short-term fix.
p.s.: I ran the script (my corrected version) and it works on master.