Using uv with PyTorch
The PyTorch ecosystem is a popular choice for deep learning research and development. You can use uv to manage PyTorch projects and PyTorch dependencies across different Python versions and environments, even controlling for the choice of accelerator (e.g., CPU-only vs. CUDA).
Installing PyTorch
From a packaging perspective, PyTorch has a few uncommon characteristics:
- Many PyTorch wheels are hosted on a dedicated index, rather than the Python Package Index (PyPI). As such, installing PyTorch typically often configuring a project to use the PyTorch index.
- PyTorch includes distinct builds for each accelerator (e.g., CPU-only, CUDA). Since there's no
standardized mechanism for specifying these accelerators when publishing or installing, PyTorch
encodes them in the local version specifier. As such, PyTorch versions will often look like
2.5.1+cpu,2.5.1+cu121, etc. - Builds for different accelerators are published to different indexes. For example, the
+cpubuilds are published on https://download.pytorch.org/whl/cpu, while the+cu121builds are published on https://download.pytorch.org/whl/cu121.
As such, the necessary packaging configuration will vary depending on both the platforms you need to support and the accelerators you want to enable.
To start, consider the following (default) configuration, which would be generated by running
uv init --python 3.12 followed by uv add torch torchvision.
In this case, PyTorch would be installed from PyPI, which hosts CPU-only wheels for Windows and macOS, and GPU-accelerated wheels on Linux (targeting CUDA 12.4):
[project]
name = "project"
version = "0.1.0"
requires-python = ">=3.12"
dependencies = [
"torch>=2.5.1",
"torchvision>=0.20.1",
]
Supported Python versions
At time of writing, PyTorch does not yet publish wheels for Python 3.13; as such projects with
requires-python = ">=3.13" may fail to resolve. See the
compatibility matrix.
This is a valid configuration for projects that want to use CPU builds on Windows and macOS, and CUDA-enabled builds on Linux. However, if you need to support different platforms or accelerators, you'll need to configure the project accordingly.
Using a PyTorch index
In some cases, you may want to use a specific PyTorch variant across all platforms. For example, you may want to use the CPU-only builds on Linux too.
In such cases, the first step is to add the relevant PyTorch index to your pyproject.toml:
We recommend the use of explicit = true to ensure that the index is only used for torch,
torchvision, and other PyTorch-related packages, as opposed to generic dependencies like jinja2,
which should continue to be sourced from the default index (PyPI).
Next, update the pyproject.toml to point torch and torchvision to the desired index:
PyTorch doesn't publish CPU-only builds for macOS, since macOS builds are always considered CPU-only.
As such, we gate on platform_system to instruct uv to ignore the PyTorch index when resolving for macOS.
PyTorch doesn't publish CUDA builds for macOS. As such, we gate on platform_system to instruct uv to ignore
the PyTorch index when resolving for macOS.
PyTorch doesn't publish CUDA builds for macOS. As such, we gate on platform_system to instruct uv to ignore
the PyTorch index when resolving for macOS.
PyTorch doesn't publish CUDA builds for macOS. As such, we gate on platform_system to instruct uv to ignore
the PyTorch index when resolving for macOS.
PyTorch doesn't publish ROCm6 builds for macOS or Windows. As such, we gate on platform_system to instruct uv to
ignore the PyTorch index when resolving for those platforms.
As a complete example, the following project would use PyTorch's CPU-only builds on all platforms:
[project]
name = "project"
version = "0.1.0"
requires-python = ">=3.12.0"
dependencies = [
"torch>=2.5.1",
"torchvision>=0.20.1",
]
[tool.uv.sources]
torch = [
{ index = "pytorch-cpu", marker = "platform_system != 'Darwin'" },
]
torchvision = [
{ index = "pytorch-cpu", marker = "platform_system != 'Darwin'" },
]
[[tool.uv.index]]
name = "pytorch-cpu"
url = "https://download.pytorch.org/whl/cpu"
explicit = true
Configuring accelerators with environment markers
In some cases, you may want to use CPU-only builds in one environment (e.g., macOS and Windows), and CUDA-enabled builds in another (e.g., Linux).
With tool.uv.sources, you can use environment markers to specify the desired index for each
platform. For example, the following configuration would use PyTorch's CPU-only builds on Windows
(and macOS, by way of falling back to PyPI), and CUDA-enabled builds on Linux:
[project]
name = "project"
version = "0.1.0"
requires-python = ">=3.12.0"
dependencies = [
"torch>=2.5.1",
"torchvision>=0.20.1",
]
[tool.uv.sources]
torch = [
{ index = "pytorch-cpu", marker = "platform_system == 'Windows'" },
{ index = "pytorch-cu124", marker = "platform_system == 'Linux'" },
]
torchvision = [
{ index = "pytorch-cpu", marker = "platform_system == 'Windows'" },
{ index = "pytorch-cu124", marker = "platform_system == 'Linux'" },
]
[[tool.uv.index]]
name = "pytorch-cpu"
url = "https://download.pytorch.org/whl/cpu"
explicit = true
[[tool.uv.index]]
name = "pytorch-cu124"
url = "https://download.pytorch.org/whl/cu124"
explicit = true
Configuring accelerators with optional dependencies
In some cases, you may want to use CPU-only builds in some cases, but CUDA-enabled builds in others,
with the choice toggled by a user-provided extra (e.g., uv sync --extra cpu vs.
uv sync --extra cu124).
With tool.uv.sources, you can use extra markers to specify the desired index for each enabled
extra. For example, the following configuration would use PyTorch's CPU-only for
uv sync --extra cpu and CUDA-enabled builds for uv sync --extra cu124:
[project]
name = "project"
version = "0.1.0"
requires-python = ">=3.12.0"
dependencies = []
[project.optional-dependencies]
cpu = [
"torch>=2.5.1",
"torchvision>=0.20.1",
]
cu124 = [
"torch>=2.5.1",
"torchvision>=0.20.1",
]
[tool.uv]
conflicts = [
[
{ extra = "cpu" },
{ extra = "cu124" },
],
]
[tool.uv.sources]
torch = [
{ index = "pytorch-cpu", extra = "cpu", marker = "platform_system != 'Darwin'" },
{ index = "pytorch-cu124", extra = "cu124" },
]
torchvision = [
{ index = "pytorch-cpu", extra = "cpu", marker = "platform_system != 'Darwin'" },
{ index = "pytorch-cu124", extra = "cu124" },
]
[[tool.uv.index]]
name = "pytorch-cpu"
url = "https://download.pytorch.org/whl/cpu"
explicit = true
[[tool.uv.index]]
name = "pytorch-cu124"
url = "https://download.pytorch.org/whl/cu124"
explicit = true
Note
Since GPU-accelerated builds aren't available on macOS, the above configuration will continue to use
the CPU-only builds on macOS via the "platform_system != 'Darwin'" marker, regardless of the extra
provided.
The uv pip interface
While the above examples are focused on uv's project interface (uv lock, uv sync, uv run,
etc.), PyTorch can also be installed via the uv pip interface.
PyTorch itself offers a dedicated interface to determine the appropriate pip command to run for a given target configuration. For example, you can install stable, CPU-only PyTorch on Linux with:
To use the same workflow with uv, replace pip3 with uv pip: