How to Switch Between CPU and GPU PyTorch Versions Using uv

Use uv sync --extra cpu or uv sync --extra cu124 to toggle between CPU-only and CUDA-enabled PyTorch builds, leveraging optional dependencies and custom package indexes defined in pyproject.toml.

The baonguyen6742/uv-install-torch repository demonstrates a robust pattern for managing multiple PyTorch variants with uv. By configuring mutually exclusive extras that point to different wheel indexes, you can seamlessly switch between CPU and GPU PyTorch versions using uv without manually editing dependency files or creating separate virtual environments.

Understanding the uv Approach to PyTorch Flavors

Unlike traditional pip workflows that require uninstalling and reinstalling PyTorch with different index URLs, uv supports optional dependencies (extras) combined with source-specific indexes. This allows the same package name (torch) to resolve to different wheels based on which extra you activate.

The Optional Dependencies Strategy

The repository defines two extras in pyproject.toml: cpu and cu124. Each extra lists the same three packages—torch, torchvision, and torchaudio—but uv resolves them from different indexes depending on which extra you specify. A conflicts rule prevents both extras from being installed simultaneously, ensuring a clean switch between variants.

Custom Package Indexes

The configuration registers two custom indexes: pytorch-cpu for CPU-only wheels and pytorch-cu124 for CUDA 12.4 wheels. These map to the official PyTorch download repositories. When you run uv sync with a specific extra, uv queries only the corresponding index, ensuring you receive the correct wheel variant without ambiguity.

Configuration in pyproject.toml

The switching mechanism is entirely defined in the project's pyproject.toml file. The relevant sections specify optional dependencies, conflict rules, source mappings, and index URLs.

[project.optional-dependencies]
cpu = ["torch==2.4.1", "torchvision==0.19.1", "torchaudio==2.4.1"]
cu124 = ["torch==2.4.1", "torchvision==0.19.1", "torchaudio==2.4.1"]

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

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

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

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

The conflicts array is critical: it tells uv that cpu and cu124 cannot coexist, forcing a clean replacement when you switch.

Switching Between CPU and GPU Builds

The workflow for switching between CPU and GPU PyTorch versions using uv requires only a single command change. Because uv manages the lock file (uv.lock) automatically, switching extras regenerates the environment to match the selected variant.

Installing the CPU-Only Build

To install the CPU-only version of PyTorch, specify the cpu extra when syncing:

uv sync --extra cpu

This command resolves torch, torchvision, and torchaudio from the pytorch-cpu index, downloading wheels compiled without CUDA support. The resulting installation is smaller and functional on machines without NVIDIA hardware.

Installing the CUDA 12.4 Build

To switch to the GPU-enabled version with CUDA 12.4 support, use the cu124 extra:

uv sync --extra cu124

Because of the conflict rule defined in pyproject.toml, uv automatically removes the CPU-only packages and replaces them with the CUDA-enabled wheels from the pytorch-cu124 index. The lock file updates to reflect the new resolution graph.

Verifying Your Installation

After syncing, verify which variant is active using a short Python command:

uv run python -c "import torch; print(torch.__version__); print('CUDA available:', torch.cuda.is_available())"

For the CPU build, the version prints as 2.4.1 and CUDA availability returns False. For the GPU build, the version includes the CUDA suffix (e.g., 2.4.1+cu124) and CUDA availability returns True if compatible hardware is present.

Handling Cache and Lock File Updates

When switching between variants, uv regenerates the uv.lock file automatically. However, if you encounter stale wheel caches or resolution errors, clean the cache explicitly:

# Remove unused cached wheels

uv cache prune

# Remove specific torch caches to force re-download

uv cache clean torch

# Re-sync with your desired extra

uv sync --extra cu124

These commands ensure that uv fetches fresh wheels from the correct index rather than reusing incompatible cached artifacts from a previous sync.

Summary

• Use extras for variants: Define cpu and cu124 extras in pyproject.toml to represent mutually exclusive PyTorch builds.
• Map sources to indexes: Use [tool.uv.sources] to bind each extra to a specific PyTorch wheel index (CPU-only or CUDA-enabled).
• Enforce conflicts: Add a [tool.uv] conflicts rule to prevent simultaneous installation of both variants.
• Switch with one command: Run uv sync --extra cpu or uv sync --extra cu124 to atomically replace the active PyTorch build.
• Verify with runtime checks: Use torch.cuda.is_available() and torch.__version__ to confirm which variant is loaded.

Frequently Asked Questions

Can I install both CPU and GPU versions simultaneously in the same environment?

No. The pyproject.toml in the baonguyen6742/uv-install-torch repository defines a conflicts rule that explicitly prevents the cpu and cu124 extras from being installed at the same time. This ensures a clean separation between variants and avoids import conflicts where PyTorch might load the wrong shared libraries.

How do I know which CUDA version (cu118, cu121, cu124) to use?

Check your NVIDIA driver version using nvidia-smi and match it to PyTorch's CUDA compatibility matrix. The repository currently implements cu124 for CUDA 12.4 support. If your system requires a different CUDA version, modify the pyproject.toml to add a new extra (e.g., cu121) and register a corresponding index pointing to https://download.pytorch.org/whl/cu121.

What happens to my existing lock file when I switch extras?

uv automatically regenerates the uv.lock file when you run uv sync with a different extra. The lock file captures the exact resolved URLs for the wheels from the selected index (CPU or CUDA). Because the extras are mutually exclusive, the new lock file replaces the previous variant's entries entirely, ensuring deterministic builds that match your current hardware requirements.

Why does uv need separate indexes for CPU and GPU wheels?

PyTorch publishes platform-specific wheels to separate repository indexes (e.g., /whl/cpu vs /whl/cu124) to reduce download size and avoid bundling CUDA libraries on CPU-only machines. By configuring distinct indexes in pyproject.toml and mapping each extra to its corresponding index via [tool.uv.sources], uv knows exactly which wheel variant to fetch without downloading unnecessary CUDA artifacts or missing required GPU libraries.

{
  "mcpServers": {
    "instagit": {
      "command": "npx",
      "args": ["-y", "instagit@latest"]
    }
  }
}