From 3d8e81478e2f9f34e367ba661ee0cd02cd12ee6b Mon Sep 17 00:00:00 2001 From: Andreas Steiner Date: Tue, 6 Feb 2024 03:27:22 -0800 Subject: [PATCH] Bumps cuda11_cudnn82 to cuda11_cudnn86. The newest `0.4.23` that is selected with `>=0.4.2` does not provide a cuda11_cudnn82 anymore: https://storage.googleapis.com/jax-releases/jax_cuda_releases.html See also https://github.com/google-research/vision_transformer/issues/295 PiperOrigin-RevId: 604593918 --- vit_jax/requirements.txt | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/vit_jax/requirements.txt b/vit_jax/requirements.txt index 35efce6..fb1cf46 100644 --- a/vit_jax/requirements.txt +++ b/vit_jax/requirements.txt @@ -6,7 +6,7 @@ einops>=0.3.0 flax>=0.6.4 git+https://github.com/google/flaxformer -jax[cuda11_cudnn82]>=0.4.2 +jax[cuda11_cudnn86]>=0.4.2 --find-links https://storage.googleapis.com/jax-releases/jax_cuda_releases.html ml-collections>=0.1.0