19 lines
593 B
Plaintext
19 lines
593 B
Plaintext
# Tensorflow.
|
|
# Cuda via pip is only on nightly right now.
|
|
# We will pin a known working version to avoid breakages (nightly breaks often).
|
|
tf-nightly[and-cuda]==2.14.0.dev20230712
|
|
|
|
# Torch.
|
|
# Pin the version used in colab currently (works with tf cuda version).
|
|
--extra-index-url https://download.pytorch.org/whl/cu118
|
|
torch==2.0.1+cu118
|
|
torchvision==0.15.2+cu118
|
|
|
|
# Jax.
|
|
# Pin the version used in colab currently (works with tf cuda version).
|
|
--find-links https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
|
|
jax[cuda11_pip]==0.4.10
|
|
|
|
# Common deps.
|
|
-r requirements-common.txt
|