diff --git a/pyproject.toml b/pyproject.toml
index ebf5dcd90..dc0a2b63b 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -237,3 +237,11 @@ quote-style = "single"
 [tool.uv]
 # Ignore uv.lock and always upgrade the package to the latest
 upgrade-package = ["jax", "jaxlib", "orbax-checkpoint"]
+
+[tool.uv.sources]
+torch = { index = "pytorch" }
+
+[[tool.uv.index]]
+name = "pytorch"
+url = "https://download.pytorch.org/whl/cpu"
+explicit = true