Thanks Sam! It seems Colab updated their version of Pytorch soon after I posted this.

The following code should fix the error.

!pip uninstall torch -y

!pip install torch==2.1