PoetryでPytorch 2.0.1をインストールするとCUDA関係の依存関係が入らない

はじめに

condaで管理していた環境をpoetryに移して、pytorchを使おうとすると……

$ python
Python 3.10.12 | packaged by conda-forge | (main, Jun 23 2023, 22:40:32) [GCC 12.3.0] on linux
Type "help", "copyright", "credits" or "license" for more information.
>>> import torch
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/opt/conda/envs/py310/lib/python3.10/site-packages/torch/__init__.py", line 229, in <module>
    from torch._C import *  # noqa: F403
ImportError: libcupti.so.11.7: cannot open shared object file: No such file or directory

pipやcondaでは正しくインストールして動作するのですが、poetryだと駄目みたいです。

原因

大方ここに書いてありました。

どうやらwheelには依存関係が正しく含まれているようですが、https://pypi.org/pypi/torch/2.0.1/json ではいくつか抜けてしまっているようです。

対処

方法1：手動で依存関係を追加

pyproject.tomlの[tool.poetry.dependencies]以下に次の依存関係を追加します。

torch = "2.0.1"
nvidia-cublas-cu11 = { version = "11.10.3.66", platform = 'linux' }
nvidia-cuda-cupti-cu11 = { version = "11.7.101", platform = 'linux' }
nvidia-cuda-nvrtc-cu11 = { version = "11.7.99", platform = 'linux' }
nvidia-cuda-runtime-cu11 = { version = "11.7.99", platform = 'linux' }
nvidia-cudnn-cu11 = { version = "8.5.0.96", platform = 'linux' }
nvidia-cufft-cu11 = { version = "10.9.0.58", platform = 'linux' }
nvidia-curand-cu11 = { version = "10.2.10.91", platform = 'linux' }
nvidia-cusolver-cu11 = { version = "11.4.0.1", platform = 'linux' }
nvidia-cusparse-cu11 = { version = "11.7.4.91", platform = 'linux' }
nvidia-nccl-cu11 = { version = "2.14.3", platform = 'linux' }
nvidia-nvtx-cu11 = { version = "11.7.91", platform = 'linux' }
triton = { version = "2.0.0", platform = 'linux' }

方法2：2.0.0を使う

2.0.0なら正しく動くようです。特に2.0.1を使う理由がなければこちらを使っても良いかと思います。
pyproject.tomlの[tool.poetry.dependencies]以下にtorch = ">=2.0.0, !=2.0.1"を追加します。

参考