alexlau commited on
Commit
279af7e
1 Parent(s): 0c4c744

set CUDA jax

Browse files
Files changed (1) hide show
  1. requirements.txt +7 -5
requirements.txt CHANGED
@@ -1,10 +1,12 @@
1
- #--find-links https://storage.googleapis.com/jax-releases/jax_releases.html
2
- #jax[cuda111]
3
- numpy>=1.16.4
4
- jax[cpu]==0.2.17
5
- jaxlib==0.1.68
 
6
  flax==0.3.4
7
  # opencv-python>=4.4.0
 
8
  Pillow>=7.2.0
9
  pyyaml>=5.3.1
10
  tensorboard>=2.4.0
 
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[cuda111]==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
10
  Pillow>=7.2.0
11
  pyyaml>=5.3.1
12
  tensorboard>=2.4.0