/poetry-torch

Installing hardware-accelerated PyTorch with Poetry on different hardware using the same `pyproject.toml`

Primary LanguagePython

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.

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

CUDA