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.