How to Install PyTorch with uv

PyTorch publishes different wheel builds for CPU, CUDA, ROCm, and XPU on separate package indexes (see Why Installing GPU Python Packages Is So Complicated for background). Getting the right build requires telling uv which index to use.

For llama.cpp models, follow Install llama-cpp-python; its backend indexes and GPU requirements differ from PyTorch’s.

Choose a build for the target machine

Start with Python 3.12 when combining PyTorch with compiled ML extensions. PyTorch’s own Python range does not guarantee that FlashAttention, xformers, or another extension has a matching wheel. Check the PyTorch install selector for the target operating system and backend.

uv add torch torchvision installs from PyPI without inspecting the GPU. Linux wheels can pull CUDA runtime dependencies even on CPU-only machines. Windows PyPI wheels are CPU-only; macOS wheels support CPU and, on supported Apple hardware, the MPS backend. Select the CPU index explicitly when a Linux environment must avoid CUDA dependencies.

CUDA wheels need a compatible NVIDIA driver on the execution host. They supply CUDA runtime libraries; a separate CUDA toolkit and nvcc are needed for source builds that compile CUDA code, not for ordinary wheel execution. The CUDA version displayed by nvidia-smi describes driver capability, not an installed toolkit.

Configure a CUDA backend in your project

In an existing uv project, define a PyTorch index in pyproject.toml and route packages to it with tool.uv.sources. This example configures CUDA 12.8 on Linux and Windows while letting macOS fall back to PyPI (since CUDA builds are not available for macOS):

[[tool.uv.index]]
name = "pytorch-cu128"
url = "https://download.pytorch.org/whl/cu128"
explicit = true

[tool.uv.sources]
torch = [
  { index = "pytorch-cu128", marker = "sys_platform == 'linux' or sys_platform == 'win32'" },
]
torchvision = [
  { index = "pytorch-cu128", marker = "sys_platform == 'linux' or sys_platform == 'win32'" },
]

Add the routed packages as direct dependencies:

uv add torch torchvision

Expect CUDA builds with +cu128 in the installation summary on Linux/Windows. uv updates pyproject.toml, uv.lock, and the project environment. Source mappings do not apply to packages that appear only as transitive dependencies.

Setting explicit = true prevents uv from searching this index for unrelated packages. Every PyTorch-related package that needs a GPU build must be listed in [tool.uv.sources]. If torchvision or torchaudio is omitted, those packages will resolve from PyPI instead of the CUDA index.

Backend indexes follow https://download.pytorch.org/whl/{backend}. An index existing does not mean it has wheels for every torch release, Python version, operating system, or GPU architecture. Use the selector to choose a supported combination; install torch, torchvision, and torchaudio together when the project needs them.

Support multiple backends with extras

Linux and Windows projects can use optional dependencies to choose a backend at install time. Replace the index/source tables from the CUDA example with this configuration:

[project.optional-dependencies]
cpu = ["torch>=2.9.1", "torchvision>=0.24.1"]
cu128 = ["torch>=2.9.1", "torchvision>=0.24.1"]

[tool.uv]
conflicts = [[{ extra = "cpu" }, { extra = "cu128" }]]

[[tool.uv.index]]
name = "pytorch-cpu"
url = "https://download.pytorch.org/whl/cpu"
explicit = true

[[tool.uv.index]]
name = "pytorch-cu128"
url = "https://download.pytorch.org/whl/cu128"
explicit = true

[tool.uv.sources]
torch = [
  { index = "pytorch-cpu", extra = "cpu" },
  { index = "pytorch-cu128", extra = "cu128" },
]
torchvision = [
  { index = "pytorch-cpu", extra = "cpu" },
  { index = "pytorch-cu128", extra = "cu128" },
]

Users then install with the extra that matches their hardware:

uv sync --extra cpu
# or
uv sync --extra cu128

The conflicts table makes the extras mutually exclusive. These CUDA extras target Linux and Windows; macOS needs its PyPI/MPS configuration instead.

Quick installs with uv pip and --torch-backend

The uv pip interface selects an index without project configuration. Create a virtual environment first; activation is unnecessary when --python names its interpreter:

uv venv --python 3.12
uv pip install --python .venv torch --torch-backend=cu128

uv reports the installed torch version with a +cu128 suffix.

Setting --torch-backend=auto makes uv detect the available GPU hardware (CUDA driver version, AMD GPU, or Intel GPU) and pick the appropriate backend. If no GPU is detected, it falls back to CPU.

uv pip install --python .venv torch --torch-backend=auto

UV_TORCH_BACKEND also accepts the backend name. The supported values depend on the installed uv; inspect uv pip install --help for the list. uv accepting a value does not establish that matching wheels exist.

Running auto on a machine with no GPU selects CPU wheels, which is the usual reason a Docker build produces a CPU-only image. How to Install CUDA PyTorch Wheels Without a GPU covers naming the target hardware instead.

Verify execution on the accelerator

On the NVIDIA execution host, save this as check_gpu.py:

import torch

print(torch.__version__, "CUDA:", torch.version.cuda, "ROCm:", torch.version.hip)
assert torch.cuda.is_available(), "Check the driver, device access, and wheel backend"
x = torch.ones((32, 32), device="cuda")
y = x @ x
assert y.is_cuda and y.sum().item() == 32768
print(torch.cuda.get_device_name(0), "matmul passed")

uv run python check_gpu.py

uv run executes the command inside the project’s virtual environment, ensuring all dependencies are installed first. The output names the GPU and ends with matmul passed; import success alone does not prove accelerator access. For the standalone uv pip environment, run uv run --no-project python check_gpu.py from its directory.

ROCm also uses the torch.cuda API; torch.version.hip distinguishes it from CUDA. Use a ROCm index and a GPU/OS combination in AMD’s support matrix. A ROCm wheel does not make an unsupported AMD card compatible. On macOS, check torch.backends.mps.is_available() and run tensors on mps instead. CPU-only execution should use device="cpu" and needs no GPU driver.

Manage PyTorch updates

Each PyTorch release supports a specific set of CUDA builds and GPU architectures. An upgrade can require a newer NVIDIA driver or matching extension wheels. Let uv’s lockfile hold the exact version.

Upgrade when you choose to, then test against the target CUDA environment:

$ uv lock --upgrade-package torch

Deferring those upgrades without also silencing PyTorch security advisories takes a Dependabot ignore rule scoped to version updates.