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 install --sync	Does not install PyTorch (import fails).
poetry install --sync -E cpu	Installs PyTorch with CPU only.
poetry install --sync -E cuda --with cuda	Installs the CUDA variant of PyTorch. Expects NVIDIA hardware.

Note

The --sync flag is optional, but useful for trying out both CPU and CUDA modes on a same machine/container, or if packages were removed. This will make sure packages that are only in the previous configuration are also removed from your current environment.

Warning

The example below is likely not what you want:

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

Embedding the choice in a script

if lspci | grep -i nvidia; then
    poetry install --sync --extras=cuda --with cuda
else
    poetry install --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