cublas Run time error

[pr4deepr]

Describe the bug
I am trying the example notebooks and successfully ran
01-data
However, when I try the training notebook and being training the model, it takes a long time to initialise and then I get the following error:
cublas runtime error : the GPU program failed to execute at C:/w/1/s/tmp_conda_3.7_044431/conda/conda-bld/pytorch_1556686009173/work/aten/src/THC/THCBlas.cu:259

Desktop (please complete the following information):

• OS: Tried this on Window 10 and Windows 11
• Graphics NVIDIA RTX 3080

Additional context
Not sure if its a compatibility issue with RTX 30 series cards.
I found a similar error for RTX 2080 cards on older pytorch
pytorch/pytorch#17334

[zacswider]

I ran into this error as well with a RTX 3070 on Windows 11. I believe the solution was installing cuda into the environment mamba install cuda -c nvidia

Also using pytorch 1.7.0, torchaudio 0.7.0, torchvision 0.8.0 and cudatoolkit 11.0.3