diff --git a/finetune/Dockerfile b/finetune/Dockerfile
index 50c1c63..a4abad2 100644
--- a/finetune/Dockerfile
+++ b/finetune/Dockerfile
@@ -5,7 +5,7 @@
 # For reference:
 # - https://docs.docker.com/build/dockerfile/frontend/#stable-channel
 ARG CUDA_VERSION=12.1.0
-FROM nvidia/cuda:${CUDA_VERSION}-cudnn8-devel-ubuntu20.04
+FROM nvidia/cuda:${CUDA_VERSION}-cudnn8-devel-ubuntu22.04
 ARG DEBIAN_FRONTEND="noninteractive"
 ENV DEBIAN_FRONTEND=${DEBIAN_FRONTEND}
 ENV MAMBA_ROOT_PREFIX=~/micromamba
@@ -85,7 +85,6 @@ ARG TORCHAUDIO_VERSION=2.3.1
 ENV TORCHAUDIO_VERSION=${TORCHAUDIO_VERSION}
 ARG PYTORCH_CUDA_VERSION=12.1
 ENV PYTORCH_CUDA_VERSION=${PYTORCH_CUDA_VERSION}
-ENV MLNX_OFED_VERSION=4.9-7.1.0.0
 ARG SETUPTOOLS_VERSION=69.5.1
 ENV SETUPTOOLS_VERSION=${SETUPTOOLS_VERSION}
 ARG USE_CUDA=1
@@ -146,6 +145,7 @@ cd ..
 rm -rf ${STAGE_DIR}/apex
 EOT
 
+ENV MLNX_OFED_VERSION=5.8-4.1.5.0
 RUN <<EOT
 #!/bin/bash
 source /opt/conda/etc/profile.d/conda.sh
@@ -161,11 +161,11 @@ python3 -m pip uninstall -y deepspeed
 # install deepspeed prepare
 # install Mellanox OFED
 mkdir -p ${STAGE_DIR}
-wget -q -O - http://www.mellanox.com/downloads/ofed/MLNX_OFED-${MLNX_OFED_VERSION}/MLNX_OFED_LINUX-${MLNX_OFED_VERSION}-ubuntu20.04-x86_64.tgz | tar xzf -
-cd MLNX_OFED_LINUX-${MLNX_OFED_VERSION}-ubuntu20.04-x86_64 
+wget -q -O - http://www.mellanox.com/downloads/ofed/MLNX_OFED-${MLNX_OFED_VERSION}/MLNX_OFED_LINUX-${MLNX_OFED_VERSION}-ubuntu22.04-x86_64.tgz | tar xzf -
+cd MLNX_OFED_LINUX-${MLNX_OFED_VERSION}-ubuntu22.04-x86_64 
 ./mlnxofedinstall --user-space-only --without-fw-update --all -q 
 cd ${STAGE_DIR} 
-rm -rf ${STAGE_DIR}/MLNX_OFED_LINUX-${MLNX_OFED_VERSION}-ubuntu20.04-x86_64*
+rm -rf ${STAGE_DIR}/MLNX_OFED_LINUX-${MLNX_OFED_VERSION}-ubuntu22.04-x86_64*
 cd ..
 # install nv_peer_mem
 rm -rf ${STAGE_DIR}
@@ -456,5 +456,5 @@ ARG NCCL_SOCKET_IFNAME='eth0'
 #     echo 'export PATH=/opt/conda/bin:$PATH' >> ~/.bashrc && \
 #     echo "source activate ${CONDA_ENV_NAME}" > ~/.bashrc
 
-# CMD ["/usr/sbin/sshd", "-D"]
-CMD ["/bin/bash", "-c", "/usr/sbin/sshd -D & while true; do sleep 1000; done"]
+CMD ["/usr/sbin/sshd", "-D"]
+# CMD ["/bin/bash", "-c", "/usr/sbin/sshd -D & while true; do sleep 1000; done"]