Align keras requirments fo torch (#2272)

keras-team · Jan 5, 2024 · 10504cb · 10504cb
1 parent dee594a
commit 10504cb
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]