jbilcke-hf HF Staff commited on
Commit
5c89d33
·
verified ·
1 Parent(s): ced4d32

Update Dockerfile

Browse files
Files changed (1) hide show
  1. Dockerfile +4 -1
Dockerfile CHANGED
@@ -68,8 +68,11 @@ RUN curl -sLo ~/miniconda.sh https://repo.continuum.io/miniconda/Miniconda3-py31
68
  && rm ~/miniconda.sh \
69
  && conda clean -ya
70
 
 
 
71
  # Make sure that Jax and cuSPARSE are properly installed for CUDA 12.x
72
- RUN conda install nvidia::libcusparse "jaxlib=*=*cuda*" jax -c conda-forge
 
73
 
74
  WORKDIR $HOME/app
75
 
 
68
  && rm ~/miniconda.sh \
69
  && conda clean -ya
70
 
71
+ ENV CONDA_OVERRIDE_CUDA="12.8"
72
+
73
  # Make sure that Jax and cuSPARSE are properly installed for CUDA 12.x
74
+ RUN conda install nvidia::cuda-toolkit nvidia::libcusparse
75
+ RUN conda install "jaxlib=*=*cuda*" jax -c conda-forge
76
 
77
  WORKDIR $HOME/app
78