Update Dockerfile

Dockerfile

@@ -31,8 +31,6 @@ RUN rm -f /etc/apt/sources.list.d/*.list && \
     # if the next line does not work, we can also try this: libegl1-mesa-dev
     libegl1 \
     libgles2 \
-    libcusparse \
-    libcusparse-dev \
     build-essential \
     libsndfile-dev \
     software-properties-common \
@@ -70,6 +68,10 @@ RUN curl -sLo ~/miniconda.sh https://repo.continuum.io/miniconda/Miniconda3-py31
  && rm ~/miniconda.sh \
  && conda clean -ya
 
+# Make sure that Jax and cuSPARSE are properly installed for CUDA 12.x
+RUN conda install nvidia::libcusparse
+RUN conda install "jaxlib=*=*cuda*" jax -c conda-forge
+
 WORKDIR $HOME/app
 
 #######################################