ampere_Dockerfile

FROM nvcr.io/nvidia/pytorch:21.09-py3

RUN pip install --upgrade pip

#install hparams
RUN pip install --upgrade git+https://github.com/Rayhane-mamah/hparams

COPY requirements.txt /tmp/requirements.txt
RUN pip install -r /tmp/requirements.txt

# Installs the wheel compatible with Cuda >= 11.1 and cudnn >= 8.0.5
RUN pip install jax==0.2.25 jaxlib==0.1.73+cuda11_cudnn82 -f https://storage.googleapis.com/jax-releases/jax_releases.html

# Installs the wheel compatible with Cuda >= 11.4 and cudnn >= 8.2
# RUN pip install jax[cuda11_cudnn82] -f https://storage.googleapis.com/jax-releases/jax_releases.html

# Install JAX extensions
RUN pip install optax==0.1.0 flax==0.3.6 trax==1.4.1
RUN pip install tensorflow==2.8.0 keras==2.8.0
# RUN pip install jax==0.2.26

ENV NUMBA_CACHE_DIR /tmp/