Solved by ptrblck in post #7 
              
              
                Thanks for the notebook! 
I can reproduce the issue in Colab using the older 1.9.0 release, while 1.9.1 works fine (also locally). 
I’ve used 
!pip install torch==1.9.1+cu111 https://download.pytorch.org/whl/torch_stable.html

in the first cell to update PyTorch.