Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
26 changes: 12 additions & 14 deletions crates/uv/src/commands/pip/compile.rs
Original file line number Diff line number Diff line change
Expand Up @@ -388,20 +388,18 @@ pub(crate) async fn pip_compile(
}

// Determine the PyTorch backend.
let torch_backend = torch_backend.map(|mode| {
if preview.is_disabled() {
warn_user!("The `--torch-backend` setting is experimental and may change without warning. Pass `--preview` to disable this warning.");
}

TorchStrategy::from_mode(
mode,
python_platform
.map(TargetTriple::platform)
.as_ref()
.unwrap_or(interpreter.platform())
.os(),
)
}).transpose()?;
let torch_backend = torch_backend
.map(|mode| {
TorchStrategy::from_mode(
mode,
python_platform
.map(TargetTriple::platform)
.as_ref()
.unwrap_or(interpreter.platform())
.os(),
)
})
.transpose()?;

// Initialize the registry client.
let client = RegistryClientBuilder::try_from(client_builder)?
Expand Down
26 changes: 12 additions & 14 deletions crates/uv/src/commands/pip/install.rs
Original file line number Diff line number Diff line change
Expand Up @@ -344,20 +344,18 @@ pub(crate) async fn pip_install(
}

// Determine the PyTorch backend.
let torch_backend = torch_backend.map(|mode| {
if preview.is_disabled() {
warn_user!("The `--torch-backend` setting is experimental and may change without warning. Pass `--preview` to disable this warning.");
}

TorchStrategy::from_mode(
mode,
python_platform
.map(TargetTriple::platform)
.as_ref()
.unwrap_or(interpreter.platform())
.os(),
)
}).transpose()?;
let torch_backend = torch_backend
.map(|mode| {
TorchStrategy::from_mode(
mode,
python_platform
.map(TargetTriple::platform)
.as_ref()
.unwrap_or(interpreter.platform())
.os(),
)
})
.transpose()?;

// Initialize the registry client.
let client = RegistryClientBuilder::try_from(client_builder)?
Expand Down
26 changes: 12 additions & 14 deletions crates/uv/src/commands/pip/sync.rs
Original file line number Diff line number Diff line change
Expand Up @@ -277,20 +277,18 @@ pub(crate) async fn pip_sync(
}

// Determine the PyTorch backend.
let torch_backend = torch_backend.map(|mode| {
if preview.is_disabled() {
warn_user!("The `--torch-backend` setting is experimental and may change without warning. Pass `--preview` to disable this warning.");
}

TorchStrategy::from_mode(
mode,
python_platform
.map(TargetTriple::platform)
.as_ref()
.unwrap_or(interpreter.platform())
.os(),
)
}).transpose()?;
let torch_backend = torch_backend
.map(|mode| {
TorchStrategy::from_mode(
mode,
python_platform
.map(TargetTriple::platform)
.as_ref()
.unwrap_or(interpreter.platform())
.os(),
)
})
.transpose()?;

// Initialize the registry client.
let client = RegistryClientBuilder::try_from(client_builder)?
Expand Down
25 changes: 16 additions & 9 deletions docs/guides/integration/pytorch.md
Original file line number Diff line number Diff line change
Expand Up @@ -433,11 +433,14 @@ $ uv pip install torch torchvision torchaudio --index-url https://download.pytor

## Automatic backend selection

In [preview](../../reference/settings.md#preview), uv can automatically select the appropriate
PyTorch index at runtime by inspecting the system configuration via `--torch-backend=auto` (or
`UV_TORCH_BACKEND=auto`):
uv supports automatic selection of the appropriate PyTorch index via the `--torch-backend=auto`
command-line argument (or the `UV_TORCH_BACKEND=auto` environment variable), as in:

```shell
$ # With a command-line argument.
$ uv pip install torch --torch-backend=auto

$ # With an environment variable.
$ UV_TORCH_BACKEND=auto uv pip install torch
```

Expand All @@ -446,12 +449,16 @@ PyTorch index for all relevant packages (e.g., `torch`, `torchvision`, etc.). If
is found, uv will fall back to the CPU-only index. uv will continue to respect existing index
configuration for any packages outside the PyTorch ecosystem.

To select a specific backend (e.g., `cu126`), set `--torch-backend=cu126` (or
`UV_TORCH_BACKEND=cu126`).
You can also select a specific backend (e.g., CUDA 12.6) with `--torch-backend=cu126` (or
`UV_TORCH_BACKEND=cu126`):

```shell
$ # With a command-line argument.
$ uv pip install torch torchvision --torch-backend=cu126

$ # With an environment variable.
$ UV_TORCH_BACKEND=cu126 uv pip install torch torchvision
```

At present, `--torch-backend` is only available in the `uv pip` interface, and only supports
detection of CUDA drivers (as opposed to other accelerators like ROCm or Intel GPUs).

As `--torch-backend` is a preview feature, it should be considered experimental and is not governed
by uv's standard [versioning policy](../../reference/policies/versioning.md). `--torch-backend` may
change or be removed entirely in future versions of uv.
Loading