add jax to pyai

2025-02-06 13:17:33 +08:00 · 2025-02-06 13:17:33 +08:00 · c1d92590b8
commit c1d92590b8
parent 4182a38b0d
2 changed files with 2 additions and 1 deletions
--- a/ScienceCompute/pyai/cpu/Dockerfile
+++ b/ScienceCompute/pyai/cpu/Dockerfile
@ -2,5 +2,6 @@ FROM eoelab.org:1027/eoeair/jupyter:py-c

 # install torch
 RUN pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cpu && \
+    pip install jax[cpu] && \
    pip cache purge

--- a/ScienceCompute/pyai/gpu/Dockerfile
+++ b/ScienceCompute/pyai/gpu/Dockerfile
@ -1,6 +1,6 @@
 FROM eoelab.org:1027/eoeair/jupyter:scipy-g

 # install torch transformers
-RUN pip install torch torchvision torchaudio && \
+RUN pip install torch torchvision torchaudio "jax[cuda12]" && \
    pip cache purge