Installing PyTorch with Poetry

Conditionally installing hardware-accelerated PyTorch with Poetry on different hardware using the same pyproject.toml can be tricky. This repo serves as a quick lookup for the configuration file and installation commands.

• Installing PyTorch with Poetry
  • Installation Modes
  • Embedding the choice in a script
  • Trying it out
  • Switching versions / upgrading
    • CUDA

Installation Modes

Command	Behavior
poetry sync	Does not install PyTorch (import fails).
poetry sync -E cpu	Installs PyTorch with CPU only.
poetry sync -E cuda --with cuda	Installs the CUDA variant of PyTorch. Expects NVIDIA hardware.

Warning

The example below is likely not what you want:

Command	Behavior
poetry sync -E cuda	Actually installs the CPU variant of PyTorch without errors or warnings.

The sync command behaves like the old poetry install --sync, and it's better suited to keep the current local state in sync with your lock file, as it will also remove dependencies missing from the lock. You probably want to use it instead of poetry install to avoid untracked outdated packages in most cases.

Embedding the choice in a script

if lspci | grep -i nvidia; then
    poetry sync --extras=cuda --with cuda
else
    poetry sync --extras=cpu
fi

Trying it out

poetry run python check-cuda.py

or

poetry run python -c "import torch; print(torch.cuda.is_available())"

Switching versions / upgrading

• Torch/Python/CUDA Compatibility Matrix of Releases
• Torchvision Versions (wheels)

CUDA

• CUDA Compute Capability numbers by GPU model
• CUDA Version Matrix by Driver Version