Dockerfile

FROM nvcr.io/nvidia/cuda:11.6.2-cudnn8-devel-ubuntu20.04

RUN apt-get update && \
    apt-get -y install python3 python3-pip && \
    apt-get clean && \
    rm -rf /var/lib/apt/lists/*

RUN python3 -m pip install jax[cuda] flax -f https://storage.googleapis.com/jax-releases/jax_releases.html
RUN python3 -m pip install tensorflow tensorflow_datasets wandb hydra-core

ENTRYPOINT [ "python3" ]

run.sh

docker run -it --rm --gpus all \
-v /etc/group:/etc/group:ro \
-v /etc/passwd:/etc/passwd:ro \
-v $PWD:/flax-example \
-v $PWD/data:/data \
-v $HOME/.netrc:$HOME/.netrc \
-v $HOME/.config/wandb:$HOME/.config/wandb \
-w /flax-example \
-u $(id -u $USER):$(id -g $USER) \
flax $@