pytorch setup with uv

CPU only:	`uv add torch --index https://download.pytorch.org/whl/cpu`
CUDA 13.0:	`uv add torch --index https://download.pytorch.org/whl/cu130`
CUDA 12.6:	`uv add torch --index https://download.pytorch.org/whl/cu126`
ROCm 7.2:	`uv add torch --index https://download.pytorch.org/whl/rocm7.2`
Auto detect:	`uv pip install torch --torch-backend=auto`
Key point:	PyTorch requires special index URLs for different compute backends
Version tags:	Backend indicated by suffix (e.g., `2.13.0+cpu`, `2.13.0+cu130`)

overview

pytorch requires special index urls for different compute backends. this guide covers uv-specific configuration for pytorch projects.

key points:

pytorch wheels hosted on separate indexes by compute type
version tags indicate backend (e.g., 2.13.0+cpu, 2.13.0+cu130)
uv handles multi-index configuration elegantly
platform-specific dependencies supported
cuda 13.0 (cu130) has been the default build since pytorch 2.11; cu128 last shipped torch 2.11, so pinned cu128 configs need to move to cu130 (or cu126 for older drivers)

quick start

cpu-only

uv add torch torchvision torchaudio --index https://download.pytorch.org/whl/cpu

cuda (latest)

uv add torch torchvision torchaudio --index https://download.pytorch.org/whl/cu130

automatic backend selection

uv pip install torch --torch-backend=auto

index urls

backend	index url	use case
cpu	`https://download.pytorch.org/whl/cpu`	no gpu acceleration
cuda 12.6	`https://download.pytorch.org/whl/cu126`	older nvidia drivers
cuda 13.0	`https://download.pytorch.org/whl/cu130`	current default for nvidia gpus
cuda 13.2	`https://download.pytorch.org/whl/cu132`	newest cuda (no torchaudio wheels)
rocm 7.2	`https://download.pytorch.org/whl/rocm7.2`	amd gpus
intel gpu	`https://download.pytorch.org/whl/xpu`	intel arc/iris

note: torchaudio is in maintenance mode and no longer versions in lockstep with torch — the current release is 2.11.0, and pytorch.org’s official install commands now include only torch torchvision. torchaudio 2.11.0 still installs fine alongside torch 2.13 on the cpu, cu126, cu130, and rocm indexes, but cu132 has no recent torchaudio wheels, so uv add torch torchvision torchaudio fails to resolve there; drop torchaudio if you use cu132.

project configuration

basic setup

# pyproject.toml
[project]
dependencies = [
    "torch>=2.13.0",
    "torchvision>=0.28.0",
    "torchaudio>=2.11.0"
]

specific index configuration

# pyproject.toml
[[tool.uv.index]]
name = "pytorch-cu130"
url = "https://download.pytorch.org/whl/cu130"

[tool.uv.sources]
torch = { index = "pytorch-cu130" }
torchvision = { index = "pytorch-cu130" }
torchaudio = { index = "pytorch-cu130" }

platform-specific setup

linux gets cuda, everything else gets cpu:

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

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

[tool.uv.sources]
torch = [
    { index = "pytorch-cpu", marker = "sys_platform != 'linux'" },
    { index = "pytorch-cu130", marker = "sys_platform == 'linux'" }
]
torchvision = [
    { index = "pytorch-cpu", marker = "sys_platform != 'linux'" },
    { index = "pytorch-cu130", marker = "sys_platform == 'linux'" }
]

optional dependencies

# pyproject.toml
[project.optional-dependencies]
cpu = [
    "torch>=2.13.0",
    "torchvision>=0.28.0",
    "torchaudio>=2.11.0"
]
cu130 = [
    "torch>=2.13.0",
    "torchvision>=0.28.0",
    "torchaudio>=2.11.0"
]

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

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

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

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

install with:

uv sync --extra cpu     # cpu-only
uv sync --extra cu130   # cuda 13.0

gpu setup

for gpu acceleration:

docker (recommended): see cuda docker setup
native: see cuda native installation

verify gpu availability:

# quick test
uv run --with torch --index https://download.pytorch.org/whl/cu130 \
    python -c "import torch; print(torch.cuda.is_available())"

common patterns

jupyter with pytorch

# temporary notebook
uv run --with jupyter --with torch --with torchvision \
    --index https://download.pytorch.org/whl/cu130 \
    jupyter lab

# permanent setup
uv add --group notebooks jupyter ipykernel matplotlib
uv add torch torchvision --index https://download.pytorch.org/whl/cu130

training script

uv can run scripts directly from urls:

# run directly from url
uv run https://michaelbommarito.com/wiki/python/scripts/check-pytorch.py

sample output with cuda:

Installed 41 packages in 186ms
pytorch version: 2.13.0+cu130
cuda available: True
cuda device: NVIDIA GeForce RTX 4070 Ti SUPER
cuda version: 13.0

cpu-only version (faster download):

# cpu version
uv run https://michaelbommarito.com/wiki/python/scripts/check-pytorch-cpu.py

sample output cpu-only:

pytorch version: 2.13.0+cpu
cuda available: False
cpu threads: 12
pytorch build: PyTorch built with:
  - GCC 13.3
  - C++ Version: 202002
  - Intel(R) oneAPI Math Kernel Library Ver...

script source:

#!/usr/bin/env -S uv run
# /// script
# dependencies = [
#   "torch",
#   "torchvision",
#   "tqdm",
#   "tensorboard"
# ]
# [tool.uv.sources]
# torch = { index = "pytorch-cu130" }
# torchvision = { index = "pytorch-cu130" }
#
# [[tool.uv.index]]
# name = "pytorch-cu130"
# url = "https://download.pytorch.org/whl/cu130"
# ///

import torch

print(f"pytorch version: {torch.__version__}")
print(f"cuda available: {torch.cuda.is_available()}")
if torch.cuda.is_available():
    print(f"cuda device: {torch.cuda.get_device_name(0)}")
    print(f"cuda version: {torch.version.cuda}")

view cpu script | view cuda script

development workflow

# init project
uv init ml-project --python 3.13
cd ml-project

# add pytorch with cuda
uv add torch torchvision torchaudio --index https://download.pytorch.org/whl/cu130

# add ml dependencies
uv add numpy pandas scikit-learn wandb
uv add --group dev pytest ipython ruff

# verify installation
uv run python -c "import torch; print(torch.cuda.is_available())"

troubleshooting

cuda not available

# check nvidia drivers
nvidia-smi

# verify cuda toolkit
nvcc --version

# reinstall with correct index
uv remove torch torchvision
uv add torch torchvision --index https://download.pytorch.org/whl/cu130

version conflicts

# clear cache
uv cache clean torch

# force reinstall
uv sync --reinstall-package torch

slow downloads

a cuda install downloads roughly 2.7gb of wheels (the torch wheel itself is ~500mb; the bundled nvidia cuda libraries make up the rest). solutions:

# use uv's built-in cache
uv cache dir  # shows cache location

# ci/cd caching
- uses: actions/cache@v4
  with:
    path: ~/.cache/uv
    key: uv-pytorch-${{ runner.os }}-${{ hashFiles('**/uv.lock') }}

# pre-download in ci
uv pip install torch --index https://download.pytorch.org/whl/cu130 --dry-run

tips

check cuda compatibility

# before installing
nvidia-smi | grep "CUDA Version"

use lock files

# lock specific versions
uv lock
# sync exact versions
uv sync --frozen

separate cpu/gpu environments

# cpu development
UV_INDEX_URL=https://download.pytorch.org/whl/cpu uv sync

# gpu training
UV_INDEX_URL=https://download.pytorch.org/whl/cu130 uv sync

specify backend explicitly

# in scripts
UV_INDEX_URL=https://download.pytorch.org/whl/cu130 uv run train.py

# or use --torch-backend flag
uv pip install torch --torch-backend=cu130