CODEWITHSUNDEEP
piprequirements.txt
MathiesW
MathiesW

Reputation: 31

Install JAX with CUDA via requirements.txt

I installed JAX[cuda] using the installation instructions from official JAX Github: pip install --upgrade "jax[cuda]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html which works nicely. Now I want to wrap this into a requirements.txt in order to save my environment.

I already tried --pre --extra-index-url https://storage.googleapis.com/jax-releases/jax_cuda_releases.html "jax[cuda]" within my requirements.txt but pip install -r requirements.txt always installs the CPU version, not the GPU version. Is there any trick to force the requirements-file to install the GPU / cuda release of JAX?

Upvotes: 0

Views: 2610

Answers (1)

sinoroc
sinoroc

Reputation: 22295

In a pip requirements.txt file, each global option should be written on its own individual line:

--pre

--find-links https://storage.googleapis.com/jax-releases/jax_cuda_releases.html

jax[cuda]

References

Upvotes: 2

Related Questions