FROM tensorflow/tensorflow:2.6.0rc2-gpu-jupyter

RUN pip install --use-feature=2020-resolver tensorflow_probability tensorflow-addons tensorflow_hub transformers datasets pandas "jax[cuda111]" -f https://storage.googleapis.com/jax-releases/jax_releases.html flax nltk spacy seaborn rich torch