How to run with 40 series cards

I have a 4090 and the code does not run with the included dockerfile. The problem is versioning conflicts between cuda and jax. Additionally Lovelace cards only support cuda 11.8 and 12 as per this [table](https://en.wikipedia.org/wiki/CUDA#GPUs_supported). Below are the changes that I made to the dockerfile to get it running.

```
(base) pcrady@pop-os:alphafold$ git diff docker/Dockerfile
diff --git a/docker/Dockerfile b/docker/Dockerfile
index 838370e..d0ae4b2 100644
--- a/docker/Dockerfile
+++ b/docker/Dockerfile
@@ -12,7 +12,7 @@
# See the License for the specific language governing permissions and
# limitations under the License.

-ARG CUDA=11.1.1
+ARG CUDA=11.8.0
FROM nvidia/cuda:${CUDA}-cudnn8-runtime-ubuntu18.04
# FROM directive resets ARGS, so we specify again (the value is retained if
# previously set).
@@ -56,11 +56,12 @@ ENV PATH="/opt/conda/bin:$PATH"
RUN conda install -qy conda==4.13.0 \
&& conda install -y -c conda-forge \
openmm=7.7.0 \
- cudatoolkit==${CUDA_VERSION} \
+ cudatoolkit=${CUDA_VERSION} \
pdbfixer \
pip \
python=3.8 \
- && conda clean --all --force-pkgs-dirs --yes
+ && conda clean --all --force-pkgs-dirs --yes \
+ && conda install cuda -c nvidia/label/cuda-11.8.0

COPY . /app/alphafold
RUN wget -q -P /app/alphafold/alphafold/common/ \
@@ -70,8 +71,8 @@ RUN wget -q -P /app/alphafold/alphafold/common/ \
RUN pip3 install --upgrade pip --no-cache-dir \
&& pip3 install -r /app/alphafold/requirements.txt --no-cache-dir \
&& pip3 install --upgrade --no-cache-dir \
- jax==0.3.25 \
- jaxlib==0.3.25+cuda11.cudnn805 \
+ jax==0.4.13 \
+ jaxlib==0.4.13+cuda11.cudnn86 \
-f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html

# Add SETUID bit to the ldconfig binary so that non-root users can run it.
```

I didnt do any validation other than folding one protein. It would be nice if the maintainers of this repo could update things to support 40 lovelace cards.