Update Cuda Reqs to align Cuda versions (#19012)

* Update index and extra url

* Align keras requirments fo torch
This commit is contained in:
Ramesh Sampath 2024-01-05 02:11:29 +05:30 committed by GitHub
parent 3553fa4914
commit 62471f289b
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 3 additions and 4 deletions

@ -1,5 +1,4 @@
# Tensorflow with cuda support.
--extra-index-url https://pypi.nvidia.com
tf-nightly[and-cuda]==2.16.0.dev20240101 # Pin a working nightly until rc0.
# Torch cpu-only version (needed for testing).

@ -2,9 +2,9 @@
tf-nightly-cpu==2.16.0.dev20240101 # 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 (needed for testing).
jax[cpu]