Skip to content

Instantly share code, notes, and snippets.

@astellon

astellon/flax.dockerfile

Created April 9, 2022 14:44
Show Gist options
  • Save astellon/b48e1ad1a8ade2c5e571a465cc9986ae to your computer and use it in GitHub Desktop.
Save astellon/b48e1ad1a8ade2c5e571a465cc9986ae to your computer and use it in GitHub Desktop.
Download ZIP
Raw
flax.dockerfile
FROM nvcr.io/nvidia/cuda:11.6.2-runtime-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
ENTRYPOINT [ "python3" ]
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment