I had the issue on my Laptop too. I had 1.9.0 installed on it. And Cuda 11.1. Python version is 3.9.
It is confusing how the same code produces different results when ran in a different way. The code in the first Colab notebook link I gave in my last reply is just segregated into separate files, like one fike for optimizer, loss and data modules (this code is in the github link) when I run it I get the error. Which is strange because both are the same code.
Also on Google Colab the torch version is 1.9.0 and CUDA 10.2 and the error still occurred.