zero_gpu doesn't support pytorch 2.2

Dockerfile

@@ -1,4 +1,4 @@
-FROM nvidia/cuda:12.3.2-cudnn9-devel-ubuntu22.04
+FROM nvidia/cuda:12.1.0-cudnn8-devel-ubuntu22.04
 ENV DEBIAN_FRONTEND=noninteractive
 RUN apt-get update && \
     apt-get upgrade -y && \
@@ -34,13 +34,13 @@ WORKDIR ${HOME}/app
 
 RUN curl https://pyenv.run | bash
 ENV PATH=${HOME}/.pyenv/shims:${HOME}/.pyenv/bin:${PATH}
-ARG PYTHON_VERSION=3.11.8
+ARG PYTHON_VERSION=3.10.11
 RUN pyenv install ${PYTHON_VERSION} && \
     pyenv global ${PYTHON_VERSION} && \
     pyenv rehash && \
     pip install --no-cache-dir -U pip setuptools wheel
 
-RUN pip install --no-cache-dir -U torch==2.2.1+cu121 torchvision>=0.16+cu121 --extra-index-url https://download.pytorch.org/whl/cu121
+RUN pip install --no-cache-dir -U torch==2.1.1+cu121 torchvision>=0.16+cu121 --extra-index-url https://download.pytorch.org/whl/cu121
 COPY --chown=1000 requirements.txt /tmp
 RUN pip install --no-cache-dir -r /tmp/requirements.txt
 

requirements.txt

@@ -1,5 +1,5 @@
-torch==2.2.1
-pytorch-lightning==2.2.0
+#torch==2.2.1
+pytorch-lightning==2.1.4
 torch-geometric==2.5.0
 pytorch-cluster==1.6.3
 scikit-learn-intelex==2024.1.0