Local inference with Blackwell (RTX 50 series) GPUs defaults to using CPU

[theendlesswatch]

This is more of a heads up than an issue, though maybe a little note in the README for other folks would be nice too.
Anyone who has an RTX 50 series GPU should activate the Python venv then run:

pip install --upgrade torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu128

This will upgrade pytorch to the latest stable version (2.8.0 as of this writing) compiled with CUDA 12.8 to enable GPU acceleration for Blackwell GPUs. I've personally tested this and everything works as expected. Though as previously noted by the author of the repo, it's currently not possible to fully utilize the GPU, save for batch inference. However, it'll still be much faster than CPU inference.

[OliBomby]

The README 'Install dependencies' step already tells you to install pytorch by following their Get Started guide. That should result in the same command you posted here. Did you follow this instruction?

[theendlesswatch]

I'm aware that the README does refer to pytorch's official website, however by default the Get Started page would direct you to install pytorch compiled with CUDA 12.6, while Blackwell specifically needs 12.8 at the minimum. Not faulting you for it, just saying that it's a bit of an "edge case" that would catch the average person off-guard.

[OliBomby]

Did you ever install CUDA manually, or was it already installed and you did not know you have 12.8 installed? Usually you have to install CUDA manually so you know which version you have.