Align keras requirments fo torch (keras-team#2272)

yuvraj-wale · Feb 8, 2024 · b91add7 · b91add7
1 parent 49c7313
commit b91add7
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 4 deletions.
diff --git a/requirements-tensorflow-cuda.txt b/requirements-tensorflow-cuda.txt
@@ -1,5 +1,4 @@
 # Tensorflow with cuda support.
---extra-index-url https://pypi.nvidia.com
 tf-nightly[and-cuda]==2.16.0.dev20231221  # Pin a working nightly until rc0.
 
 # Torch cpu-only version.

diff --git a/requirements-torch-cuda.txt b/requirements-torch-cuda.txt
@@ -2,9 +2,9 @@
 tf-nightly-cpu==2.16.0.dev20231221  # Pin a working nightly until rc0.
 
 # Torch with cuda support.
---extra-index-url https://download.pytorch.org/whl/cu118
-torch==2.1.2+cu118
-torchvision==0.16.2+cu118
+--extra-index-url https://download.pytorch.org/whl/cu121
+torch==2.1.2
+torchvision==0.16.2
 
 # Jax cpu-only version.
 jax[cpu]