Set torch url index in pyproject.toml (#16802)

2026-02-01 13:23:52 +08:00
parent 429ef988bc
commit c7d53fa26a
2 changed files with 22 additions and 6 deletions
--- a/python/pyproject.toml
+++ b/python/pyproject.toml
@@ -82,6 +82,22 @@ dependencies = [
  "grpcio-health-checking==1.75.1", # required for Kubernetes gRPC health probes
 ]

+[[tool.uv.index]]
+name = "pypi"
+url = "https://pypi.org/simple"
+default = true
+
+[[tool.uv.index]]
+name = "torch-cu129"
+url = "https://download.pytorch.org/whl/cu129"
+explicit = true
+
+[tool.uv.sources]
+torch = [
+  { index = "pypi", marker = "platform_machine == 'x86_64'"},
+  { index = "torch-cu129", marker = "platform_machine == 'aarch64'"},
+]
+
 [project.optional-dependencies]
 checkpoint-engine = ["checkpoint-engine==0.1.2"]
 diffusion = [