Specific jaxlib version?

Start with the colabfold_batch entry point and reproduce the reported command after the listed pip and conda installation steps. Compare the installed JAX 0.3.13, jaxlib 0.3.10, CUDA 11.1.1, and cuDNN 8.4.1.50 versions with the command-line behavior; done means the documented local installation produces a functioning GPU workflow without the unrecognized-argument error.

Expected Behavior

Fully functioning colab_batch with GPU.

Current Behavior

colabfold_batch: error: unrecognized arguments: --use-gpu-relax

Steps to Reproduce (for bugs)

Install as described for local installation:
./pip install --upgrade pip
./pip install "colabfold[alphafold] @ git+https://github.com/sokrypton/ColabFold"
./pip install "jax[cuda111]>=0.3.8,<0.4" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
./conda install -y cudatoolkit=11.1.1 cudnn=8.4.1.50
./conda install -y -c conda-forge -c bioconda kalign2=2.04 hhsuite=3.3.0
./conda install -y -c conda-forge openmm=7.5.1 pdbfixer

ColabFold Output (for bugs)

usage: colabfold_batch [-h] [--stop-at-score STOP_AT_SCORE]
[--num-recycle NUM_RECYCLE] [--num-models {1,2,3,4,5}]
[--recompile-padding RECOMPILE_PADDING]
[--model-order MODEL_ORDER] [--host-url HOST_URL]
[--data DATA]
[--msa-mode {MMseqs2 UniRef+Environmental),MMseqs2 (UniRef only,single_sequence}]
[--model-type {auto,AlphaFold2-ptm,AlphaFold2-multimer}]
[--amber] [--templates] [--env] [--cpu]
[--rank {auto,plddt,ptmscore,multimer}]
[--pair-mode {unpaired,paired,unpaired+paired}]
[--recompile-all-models]
[--sort-queries-by {none,length,random}]
[--save-single-representations]
[--save-pair-representations] [--zip]
[--overwrite-existing-results]
input results
colabfold_batch: error: unrecognized arguments: --use-gpu-relax

Context

Resulting installation contains:

cudatoolkit 11.1.1 ha002fc5_10 conda-forge
cudnn 8.4.1.50 hed8a83a_0 conda-forge
jax 0.3.13
jaxlib 0.3.10

Your Environment

Local install
CentOS Linux release 7.9.2009 (Core)
NVIDIA-SMI 460.39 Driver Version: 460.39 CUDA Version: 11.2
GeForce GTX 108