Updated dockerfile to match jaxlib version

KosinskiLab · Feb 20, 2024 · 89fcb95 · 89fcb95
1 parent 6cd2d55
commit 89fcb95
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 6 deletions.
diff --git a/workflow/docker/pulldown.dockerfile b/workflow/docker/pulldown.dockerfile
@@ -50,8 +50,7 @@ RUN wget -q -P /tmp \
 
 ENV PATH="/opt/conda/bin:$PATH"
 ENV LD_LIBRARY_PATH="/opt/conda/lib:$LD_LIBRARY_PATH"
-RUN conda install -qy conda==23.5.2 \
-    && conda install -y -c conda-forge -c bioconda \
+RUN conda install -y -c conda-forge -c bioconda \
       openmm==8.0 \
       cudatoolkit==${CUDA_VERSION} \
       pdbfixer==1.9 \
@@ -70,7 +69,7 @@ RUN pip3 install --upgrade pip --no-cache-dir \
     && pip3 install --upgrade --no-cache-dir \
       alphapulldown==1.0.4 \
       pytest \
-      "jax[cuda11_pip]" \
+      jax==0.4.23 \
       jaxlib==0.4.23 \
       -f https://storage.googleapis.com/jax-releases/jaxlib_nightly_cuda_releases.html
 

diff --git a/workflow/source/alphafold_inference.py b/workflow/source/alphafold_inference.py
@@ -198,9 +198,7 @@ def predict_multimer(
         flags_dict["gradient_msa_depth"] = gradient_msa_depth
         flags_dict["model_names_custom"] = model_names
         flags_dict["msa_depth"] = msa_depth
-
-
-    if not isinstance(multimer, MultimericObject):
+    else:
         multimer.input_seqs = [multimer.sequence]
 
     model_runners, random_seed = create_model_runners_and_random_seed(**flags_dict)