diff --git a/tox.ini b/tox.ini
index b16c0d2f0b..cfd0beedfd 100644
--- a/tox.ini
+++ b/tox.ini
@@ -83,8 +83,8 @@ passenv =
 commands =
     python -c "import os; os.system('install-custom-pkgs --install-boto-wheels')"
     pip install 'apache-airflow==2.9.3' --constraint "https://raw.githubusercontent.com/apache/airflow/constraints-2.9.3/constraints-3.8.txt"
-    pip install 'torch==2.0.1+cpu' -f 'https://download.pytorch.org/whl/torch_stable.html'
-    pip install 'torchvision==0.15.2+cpu' -f 'https://download.pytorch.org/whl/torch_stable.html'
+    pip install 'torch==2.0.1' -f 'https://download.pytorch.org/whl/torch_stable.html'
+    pip install 'torchvision==0.15.2' -f 'https://download.pytorch.org/whl/torch_stable.html'
     pip install 'dill>=0.3.8'
 
     pytest {posargs}