UV安装GPU版本PyTorch

经过同事推荐，开始尝试使用uv管理Python环境，效果相当不错。

安装PyTorch遇到的问题

但在安装PyTorch时，采用默认的uv add方式会报错，而使用uv pip install安装PyTorch的cuda版本，虽然没有问题，但并不能同步到pyproject.toml文件中。

阅读PyTorch官网文档，我本机cuda环境是12.6，使用pip安装，其命令为：

pip3 install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu126

那么，对应的uv pip 命令是：

uv pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu126

解决方案

如前所述，这样安装并不便于项目包的统一管理，根据UV官网的PyTorch集成方法，我们可以通过修改pyproject.toml文件然后uv sync的方式实现GPU版本PyTorch的安装。

根据官网文档，可以在pyproject.toml文件中添加以下内容：

...
dependencies = [
    "torch>=2.6.0",
    "torchaudio>=2.6.0",
    "torchvision>=0.21.0",
]

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

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

注：因为我本机的CUDA版本是12.6，所以修改内容的后缀是-cu126，官网示例最新为-cu124，对应cuda 12.4，实际上我这个版本也是可用的，只是指向问题。

亲测有效，问题解决。

import torch

print(torch.__version__)
print(torch.cuda.is_available())

2.6.0+cu126
True