File tree Expand file tree Collapse file tree 1 file changed +4
-8
lines changed Expand file tree Collapse file tree 1 file changed +4
-8
lines changed Original file line number Diff line number Diff line change @@ -253,18 +253,14 @@ simply run:
253253conda install jax -c conda-forge
254254```
255255
256- To install it on a machine with an NVIDIA GPU, run:
256+ If you run this command on machine with an NVIDIA GPU, this should install a CUDA-enabled package of ` jaxlib ` .
257+
258+ To ensure that the jax version you are installing is indeed CUDA-enabled, run:
257259
258260``` bash
259- conda install " jaxlib=*=*cuda*" jax cuda-nvcc - c conda-forge -c nvidia
261+ conda install " jaxlib=*=*cuda*" jax - c conda-forge
260262```
261263
262- Note the ` cudatoolkit ` distributed by ` conda-forge ` is missing ` ptxas ` , which
263- JAX requires. You must therefore either install the ` cuda-nvcc ` package from
264- the ` nvidia ` channel, or install CUDA on your machine separately so that ` ptxas `
265- is in your path. The channel order above is important (` conda-forge ` before
266- ` nvidia ` ).
267-
268264If you would like to override which release of CUDA is used by JAX, or to
269265install the CUDA build on a machine without GPUs, follow the instructions in the
270266[ Tips & tricks] ( https://conda-forge.org/docs/user/tipsandtricks.html#installing-cuda-enabled-packages-like-tensorflow-and-pytorch )
You can’t perform that action at this time.
0 commit comments