如何告诉 PyTorch 使用哪个 CUDA 版本？

Question

我的 Ubuntu 16.04 机器上安装了两个版本的 CUDA：9.0 和 10.1。 它们分别位于 /usr/local/cuda-9.0 和 /usr/local/10.1。 如果我通过 pip (pip install torch==1.6.0) 安装 PyTorch 1.6.0（需要 CUDA 10.1），它使用 9.0 版，因此没有检测到 GPU。我已将 LD_LIBRARY_PATH 更改为 "/usr/local/cuda-10.1/lib64:/usr/local/cuda-10.1/cuda/extras/CUPTI/lib64"，但 PyTorch 仍在使用 CUDA 9.0。 如何告诉 PyTorch 使用 CUDA 10.1？

Answer 1

使用不同版本的 CUDA 构建的手电筒的预制轮子可在 torch stable releases page 获得。例如，您可以像这样安装使用 CUDA v11.1 构建的 torch v1.9.0：

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

但并非所有组合都可用。