Skip to content

ColabFold runs on CPU during evaluation due to missing JAX CUDA support #5

@sven-lange

Description

@sven-lange

During the evaluation step, ColabFold reports "WARNING: no GPU detected, will be using CPU" despite a GPU being allocated via SLURM. This makes evaluation extremely slow.

Root cause: The genie3 conda environment installs jaxlib without CUDA support.

Fix: pip install "jax[cuda]" "openmm[cuda12]"

Note: pip will warn about a nvidia-cudnn-cu12 version conflict with PyTorch, but both PyTorch and JAX work correctly at runtime despite this warning.

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type
    No fields configured for issues without a type.

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions