dawnbench_jax/Dockerfile
2024-11-11 23:54:14 -07:00

18 lines
309 B
Docker

FROM docker.io/python:3.12.7-bookworm
WORKDIR /workspace
RUN apt-get update && apt-get install -y \
build-essential \
git \
curl \
&& rm -rf /var/lib/apt/lists/*
RUN pip install --no-cache-dir \
jax[cuda12] \
equinox \
torch \
optax \
torchvision
CMD ["/usr/bin/bash"]