chore: Pytorch gpu

alexandrainst · Oct 31, 2024 · 8eceb45 · 8eceb45
1 parent c5b797a
commit 8eceb45
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 4 deletions.
diff --git a/pyproject.toml b/pyproject.toml
@@ -59,16 +59,19 @@ dev-dependencies = [
 ]
 
 [tool.uv.sources]
-torch = { index = "pytorch-cpu", marker = "extra == 'onprem_cpu'" }
+torch = [
+    { index = "pytorch-cpu", marker = "extra == 'onprem_cpu'" },
+    { index = "pytorch-cu126", marker = "extra != 'onprem_cpu'" },
+]
 
 [[tool.uv.index]]
 name = "pytorch-cpu"
 url = "https://download.pytorch.org/whl/cpu"
 explicit = true
 
 [[tool.uv.index]]
-name = "pypi"
-url = "https://pypi.org/simple"
+name = "pytorch-cu126"
+url = "https://download.pytorch.org/whl/cu126"
 
 [project.scripts]
 ragger-demo = "ragger.cli:run_demo"

diff --git a/uv.lock b/uv.lock