requirements_gpu.txt

transformers==4.19.2
absl-py==0.13.0
biopython==1.79
chex==0.0.7
dm-haiku==0.0.5
dm-tree==0.1.6
docker==5.0.0
immutabledict==2.0.0
jax[cuda]<0.3.0
jaxlib==0.1.76
ml-collections==0.1.0
numpy==1.19.5
pandas==1.3.4
scipy==1.7.0
tensorflow-gpu==2.5.0
torch
plotly
GPUtil
numba
ray
protobuf<4
-f https://storage.googleapis.com/jax-releases/jax_releases.html