GPU not found on WSL

[wolfgangmeyers]

Getting the following output when trying to run on WSL. I've followed all of the steps that should make it work, pytorch seems to think there is GPU support? Not sure if this is a bogus error or not.

(.venv) wolfgang@DESKTOP-4RO0FD7:~/code$ python
Python 3.8.10 (default, Mar 15 2022, 12:22:08)
[GCC 9.4.0] on linux
Type "help", "copyright", "credits" or "license" for more information.
>>> import torch
>>> torch.cuda.device_count()
1
>>> exit()
(.venv) wolfgang@DESKTOP-4RO0FD7:~/code$ cd dalle-playground/backend/
(.venv) wolfgang@DESKTOP-4RO0FD7:~/code/dalle-playground/backend$ python app.py --port 8080 --model_version mega_full
--> Starting DALL-E Server. This might take up to two minutes.
WARNING:absl:No GPU/TPU found, falling back to CPU. (Set TF_CPP_MIN_LOG_LEVEL=0 and rerun for more info.)

[wolfgangmeyers]

Yeah, looks like it's running on the cpu. Might dig further to see if I can find out why.

[wolfgangmeyers]

Seems related to jax-ml/jax#5231

[wolfgangmeyers]

Finally managed to get it to work. Needed to do the following:

Install cuda: https://developer.nvidia.com/cuda-downloads?target_os=Linux&target_arch=x86_64&Distribution=WSL-Ubuntu&target_version=2.0&target_type=deb_local

Install cuDNN: https://developer.nvidia.com/rdp/cudnn-download
This threw me for a bit. It drops a bunch of .deb files in /var/cudnn-local-repo-ubuntu2004-8.4.1.50. I needed to install a .deb from that folder.

Then I followed installation instructions for jax here: https://github.com/google/jax#installation

[matthew-lowe]

Did you get any GPU usage when running? I had the same issue, was able to fix it (error went away, both jax and pytorch identified my GPU and said it was available) but when running the model there was no GPU usage, CPU usage maxed out on a single core and it took as long as before

[matthew-lowe]

Well I think I can conclude that my GPU usage (at least in task manager) is a bit scuffed, since playing RDR2 doesn't show any considerable increase in usage