Update Cuda Reqs to align Cuda versions (#19012)
* Update index and extra url * Align keras requirments fo torch
This commit is contained in:
parent
3553fa4914
commit
62471f289b
@ -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]
|
||||
|
Loading…
Reference in New Issue
Block a user