infochip
infochip

Reputation: 223

How to tell PyTorch which CUDA version to take?

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?

Upvotes: 9

Views: 29766

Answers (1)

Mehraban
Mehraban

Reputation: 3324

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.

Upvotes: 9

Related Questions