juancopi81 commited on
Commit
894d943
1 Parent(s): 85250f0

Change Jax version

Browse files
Files changed (1) hide show
  1. requirements.txt +3 -1
requirements.txt CHANGED
@@ -1,7 +1,9 @@
1
  tensorflow
2
  nest-asyncio
3
  pyfluidsynth==1.3.0
4
- jax[cpu]
 
 
5
  # pin CLU for python 3.7 compatibility
6
  clu==0.0.7
7
  # pin Orbax to use Checkpointer
 
1
  tensorflow
2
  nest-asyncio
3
  pyfluidsynth==1.3.0
4
+ # downgrade jax
5
+ # TODO(iansimon): remove this once we can use the latest t5x
6
+ jax[cpu]==0.3.15 -f https://storage.googleapis.com/jax-releases/jax_releases.html
7
  # pin CLU for python 3.7 compatibility
8
  clu==0.0.7
9
  # pin Orbax to use Checkpointer