alexlau commited on
Commit
6d2eea0
1 Parent(s): 387e648

cpu jax installation

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