Skip to content

Instantly share code, notes, and snippets.

@paul-english
Created April 15, 2026 02:24
Show Gist options
  • Select an option

  • Save paul-english/64f8da7efc3a695131df9242b5b4ea14 to your computer and use it in GitHub Desktop.

Select an option

Save paul-english/64f8da7efc3a695131df9242b5b4ea14 to your computer and use it in GitHub Desktop.
Multi gpu torch pyproject
[project]
name = "some-torch-project"
version = "0.1.0"
description = "Add your description here"
readme = "README.md"
requires-python = ">=3.12"
dependencies = [
"torchmetrics>=1.8.0",
]
[project.optional-dependencies]
# Pick exactly one — mutually exclusive (see [tool.uv.conflicts]). Example: uv sync --extra cu128
noindex = [
"torch>=2.9.1",
"torchvision>=0.24.1",
]
cpu = [
"torch>=2.9.1",
"torchvision>=0.24.1",
]
cu128 = [
"torch>=2.9.1",
"torchvision>=0.24.1",
]
rocm72 = [
"torch>=2.9.1",
"torchvision>=0.24.1",
# Matches torch's dependency on the ROCm 7.2 index (see lock resolution).
"triton-rocm==3.6.0 ; sys_platform == 'linux'",
]
[build-system]
requires = ["hatchling"]
build-backend = "hatchling.build"
# PyTorch wheels: https://docs.astral.sh/uv/guides/integration/pytorch/
[tool.uv]
conflicts = [
[
{ extra = "noindex" },
{ extra = "cpu" },
{ extra = "cu128" },
{ extra = "rocm72" },
],
]
[tool.uv.sources]
torch = [
{ index = "pytorch-cpu", extra = "cpu" },
{ index = "pytorch-cu128", extra = "cu128", marker = "sys_platform == 'linux' or sys_platform == 'win32'" },
{ index = "pytorch-rocm72", extra = "rocm72", marker = "sys_platform == 'linux'" },
]
torchvision = [
{ index = "pytorch-cpu", extra = "cpu" },
{ index = "pytorch-cu128", extra = "cu128", marker = "sys_platform == 'linux' or sys_platform == 'win32'" },
{ index = "pytorch-rocm72", extra = "rocm72", marker = "sys_platform == 'linux'" },
]
triton-rocm = [
{ index = "pytorch-rocm72", extra = "rocm72", marker = "sys_platform == 'linux'" },
]
[[tool.uv.index]]
name = "pytorch-cpu"
url = "https://download.pytorch.org/whl/cpu"
explicit = true
[[tool.uv.index]]
name = "pytorch-cu128"
url = "https://download.pytorch.org/whl/cu128"
explicit = true
[[tool.uv.index]]
name = "pytorch-rocm72"
url = "https://download.pytorch.org/whl/rocm7.2"
explicit = true
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment