jax[cuda] working

cai4cai · Oct 21, 2024 · b4f2e71 · b4f2e71
1 parent cc66ff9
commit b4f2e71
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 2 deletions.
diff --git a/.devcontainer/Dockerfile b/.devcontainer/Dockerfile
@@ -30,7 +30,7 @@ ARG CUDATOOLKIT_VERSION=12.4
 RUN conda install pytorch=${PYTORCH_VERSION} pytorch-cuda=${CUDATOOLKIT_VERSION} -c pytorch -c nvidia
 
 # Handle environment.yml if it exists
-RUN echo env_change_20241021
+RUN echo env_change_20241021_2
 COPY environment.yml* noop.txt /tmp/conda-tmp/
 RUN if [ -f "/tmp/conda-tmp/environment.yml" ]; then \
     /opt/conda/bin/conda env update -n base -f /tmp/conda-tmp/environment.yml; \

diff --git a/.devcontainer/environment.yml b/.devcontainer/environment.yml
@@ -5,7 +5,6 @@ channels:
 dependencies:
   - numpy
   - cupy
-  - jax[cuda]
   - scipy
   - pre-commit==3.7.1
   - black==24.4.2
@@ -18,3 +17,6 @@ dependencies:
   - libmamba
   - libmambapy
   - libarchive
+  - pip
+  - pip:
+    - "jax[cuda12]"