simonduerr commited on
Commit
17155bc
·
1 Parent(s): 8fc30a2

Update requirements.txt

Browse files
Files changed (1) hide show
  1. requirements.txt +2 -2
requirements.txt CHANGED
@@ -6,13 +6,13 @@ dm-haiku==0.0.5
6
  dm-tree==0.1.6
7
  docker==5.0.0
8
  immutabledict==2.0.0
9
- jax<0.3.0
10
  jaxlib==0.1.76
11
  ml-collections==0.1.0
12
  numpy==1.19.5
13
  pandas==1.3.4
14
  scipy==1.7.0
15
- tensorflow-cpu==2.5.0
16
  torch
17
  plotly
18
  GPUtil
 
6
  dm-tree==0.1.6
7
  docker==5.0.0
8
  immutabledict==2.0.0
9
+ jax[cuda]<0.3.0
10
  jaxlib==0.1.76
11
  ml-collections==0.1.0
12
  numpy==1.19.5
13
  pandas==1.3.4
14
  scipy==1.7.0
15
+ tensorflow-gpu==2.5.0
16
  torch
17
  plotly
18
  GPUtil