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

Question

我的 Ubuntu 16.04 机器上安装了两个版本的 CUDA：9.0 和 10.1。它们分别位于/usr/local/cuda-9.0和/usr/local/10.1。如果我通过 pip ( ) 安装 PyTorch 1.6.0（需要 CUDA 10.1）pip install torch==1.6.0，它使用版本 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 的预制轮可在torch 稳定版本页面上找到。例如，您可以安装使用 CUDA v11.1 构建的 torch v1.9.0，如下所示：

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

但并非所有组合都可用。