This is a bug in the autodiff, I would recommend to file an issue on the PyTorch github (and crosslink here and the issue). As someone who sometimes looks into PyTorch issues, thank you for making a reproducing example. These are gold to anyone trying to fix things!
Best regards
Thomas