How to tell PyTorch which CUDA version to take?

user_13600920 0
Viewed 6425

I have two version of CUDA installed on my Ubuntu 16.04 machine: 9.0 and 10.1. They are located in /usr/local/cuda-9.0 and /usr/local/10.1 respectively. If I install PyTorch 1.6.0 (which needs CUDA 10.1) via pip (pip install torch==1.6.0), it uses version 9.0 and thus detects no GPUs. I already changed my LD_LIBRARY_PATH to "/usr/local/cuda-10.1/lib64:/usr/local/cuda-10.1/cuda/extras/CUPTI/lib64" but PyTorch is still using CUDA 9.0. How do I tell PyTorch to use CUDA 10.1?

pytorch
1 Answers

Prebuilt wheels for torch built with different versions of CUDA are available at torch stable releases page. For example you can install torch v1.9.0 built with CUDA v11.1 like this:

pip install --upgrade torch==1.9.0+cu111 -f https://download.pytorch.org/whl/torch_stable.html

But not all the combinations are available.

user_18057560
Related