diff --git a/pyproject.toml b/pyproject.toml
index f410a72c..759ac167 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -20,10 +20,9 @@ numpyro = {version = "*", optional = true}
 arviz = {version = "*", optional = true}
 funsor = {version = "*", optional = true}
 torch = [
-     {version = "<=2.1.0", platform = "darwin"},
-     {version = "<=2.1.0", platform = "linux", source = "torch"},
-     {version = "<=2.1.0", platform = "win32", source = "torch"},
-]
+     {version = "^2.0.1", platform = "darwin"},
+     {version = "^2.0.1", platform = "linux", source = "torch"},
+     {version = "^2.0.1", platform = "win32", source = "torch"}, ]
 lightning = "*"
 
 [[tool.poetry.source]]
@@ -46,7 +45,7 @@ opentsne = "*"
 umap-learn = "*"
 rdata = "*"
 torch = "*"
-torchvision = "*"
+torchvision = "0.17.0"
 pre-commit = "*"