Alex Lau commited on
Commit
6c4ca29
1 Parent(s): 492ee73

install cuda version of jax and jaxlib

Browse files
Files changed (1) hide show
  1. requirements.txt +6 -6
requirements.txt CHANGED
@@ -1,11 +1,11 @@
1
  # in CUDA mode, can set src.config.py::DOWNSAMPLE=2
2
- #jaxlib==0.1.68
3
- #--find-links https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
4
- #jax[cuda112]==0.2.17
5
  # in CPU mode, set src.config.py::DOWNSAMPLE=4
6
- --find-links https://storage.googleapis.com/jax-releases/jax_releases.html
7
- jax[cpu]==0.2.17
8
- jaxlib==0.1.68
9
  flax==0.3.4
10
  # opencv-python>=4.4.0
11
  numpy>=1.16.4
1
  # in CUDA mode, can set src.config.py::DOWNSAMPLE=2
2
+ --find-links https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
3
+ jax[cuda112]==0.2.17
4
+ jaxlib==0.1.68+cuda111
5
  # in CPU mode, set src.config.py::DOWNSAMPLE=4
6
+ #--find-links https://storage.googleapis.com/jax-releases/jax_releases.html
7
+ #jax[cpu]==0.2.17
8
+ #jaxlib==0.1.68
9
  flax==0.3.4
10
  # opencv-python>=4.4.0
11
  numpy>=1.16.4