Revision 4d8bce1b8515c68a8432d379d964ea3e6bea463c authored by Yash Katariya on 19 October 2021, 02:55:51 UTC, committed by jax authors on 19 October 2021, 02:56:22 UTC
```
# Installs Cuda 11 with Cudnn 8.2
$ pip install jax[cuda] -f https://storage.googleapis.com/jax-releases/jax_releases.html

$ pip install jax[cuda=11,cudnn=82] -f https://storage.googleapis.com/jax-releases/jax_releases.html

$ pip install jax[cuda=11,cudnn=805] -f https://storage.googleapis.com/jax-releases/jax_releases.html
```

PiperOrigin-RevId: 404134291
1 parent 6bd0c72
History
File Mode Size
api_benchmark.py -rw-r--r-- 7.4 KB
benchmark.py -rw-r--r-- 5.6 KB
pmap_benchmark.py -rw-r--r-- 4.9 KB

back to top