diff --git a/docker/pytorch/training/1.13.1/Dockerfile.neuronx b/docker/jax/training/0.4/Dockerfile.neuronx
similarity index 60%
rename from docker/pytorch/training/1.13.1/Dockerfile.neuronx
rename to docker/jax/training/0.4/Dockerfile.neuronx
index 378747a..99f213d 100644
--- a/docker/pytorch/training/1.13.1/Dockerfile.neuronx
+++ b/docker/jax/training/0.4/Dockerfile.neuronx
@@ -1,23 +1,21 @@
-FROM public.ecr.aws/docker/library/ubuntu:20.04
+FROM public.ecr.aws/docker/library/ubuntu:22.04
 
-LABEL maintainer="Amazon AI"
 LABEL dlc_major_version="1"
+LABEL maintainer="Amazon AI"
 
 # Neuron SDK components version numbers
-ARG NEURONX_FRAMEWORK_VERSION=1.13.1.1.16.0
-ARG NEURONX_DISTRIBUTED_VERSION=0.9.0
-ARG NEURONX_DISTRIBUTED_TRAINING_VERSION=1.0.1
-ARG NEURONX_CC_VERSION=2.15.143.0
-ARG NEURONX_COLLECTIVES_LIB_VERSION=2.22.33.0-d2128d1aa
-ARG NEURONX_RUNTIME_LIB_VERSION=2.22.19.0-5856c0b42
-ARG NEURONX_TOOLS_VERSION=2.19.0.0
+ARG NEURONX_RUNTIME_LIB_VERSION=2.23.110.0-9b5179492
+ARG NEURONX_COLLECTIVES_LIB_VERSION=2.23.133.0-3e70920f2
+ARG NEURONX_TOOLS_VERSION=2.20.204.0
+ARG NEURONX_CC_VERSION=2.16.345.0
+ARG NEURONX_JAX_TRAINING_VERSION=0.1.2
 
 ARG PYTHON=python3.10
 ARG PYTHON_VERSION=3.10.12
 ARG PIP=pip3
 ARG OMPI_VERSION=4.1.5
 
-# This arg required to stop docker build waiting for region configuration while installing tz data from ubuntu 20
+# This arg required to stop docker build waiting for region configuration while installing tz data from ubuntu 22
 ARG DEBIAN_FRONTEND=noninteractive
 
 # Python won’t try to write .pyc or .pyo files on the import of source modules
@@ -32,9 +30,6 @@ ENV LD_LIBRARY_PATH="${LD_LIBRARY_PATH}:/opt/amazon/efa/lib"
 ENV LD_LIBRARY_PATH="${LD_LIBRARY_PATH}:/opt/amazon/efa/lib64"
 ENV LD_LIBRARY_PATH="${LD_LIBRARY_PATH}:/opt/amazon/openmpi/lib64"
 ENV LD_LIBRARY_PATH="${LD_LIBRARY_PATH}:/usr/local/lib"
-ENV PATH /opt/aws/neuron/bin/:$PATH
-ENV SAGEMAKER_TRAINING_MODULE=sagemaker_pytorch_container.training:main
-ENV DGLBACKEND=pytorch
 
 RUN apt-get update \
  && apt-get upgrade -y \
@@ -45,45 +40,35 @@ RUN apt-get update \
     curl \
     emacs \
     git \
+    gnupg2 \
+    gpg-agent \
     jq \
     libopencv-dev \
-    openjdk-8-jdk-headless \
-    openjdk-8-jdk \
-    openjdk-8-jre \
     libglib2.0-0 \
     libgl1-mesa-glx \
     libsm6 \
     libxext6 \
     libxrender-dev \
-    openjdk-11-jdk \
-    software-properties-common \
-    wget \
-    unzip \
-    vim \
-    zlib1g-dev \
-    openssl \
     libssl-dev \
     libsqlite3-dev \
     libgdbm-dev \
     libc6-dev \
     libbz2-dev \
     libncurses-dev \
-    tk-dev \
     libffi-dev \
     libcap-dev \
-    gnupg2 \
-    gpg-agent \
- && rm -rf /var/lib/apt/lists/* \
- && apt-get clean
-
-RUN echo "deb https://apt.repos.neuron.amazonaws.com focal main" > /etc/apt/sources.list.d/neuron.list
-RUN wget -qO - https://apt.repos.neuron.amazonaws.com/GPG-PUB-KEY-AMAZON-AWS-NEURON.PUB | apt-key add -
-
-RUN apt-get update \
- && apt-get install -y \
-    aws-neuronx-tools=$NEURONX_TOOLS_VERSION \
-    aws-neuronx-collectives=$NEURONX_COLLECTIVES_LIB_VERSION \
-    aws-neuronx-runtime-lib=$NEURONX_RUNTIME_LIB_VERSION \
+    libhwloc-dev \
+    openjdk-8-jdk-headless \
+    openjdk-8-jdk \
+    openjdk-8-jre \
+    openjdk-11-jdk \
+    openssl \
+    software-properties-common \
+    tk-dev \
+    unzip \
+    wget \
+    vim \
+    zlib1g-dev \
  && rm -rf /var/lib/apt/lists/* \
  && rm -rf /tmp/tmp* \
  && apt-get clean
@@ -100,6 +85,15 @@ RUN mkdir -p /tmp/openmpi \
  && ldconfig \
  && rm -rf /tmp/openmpi
 
+# Install packages and configure SSH for MPI operator in k8s
+RUN apt-get update && apt-get install -y openmpi-bin openssh-server \
+ && mkdir -p /var/run/sshd \
+ && echo "    UserKnownHostsFile /dev/null" >> /etc/ssh/ssh_config \
+ && echo "    StrictHostKeyChecking no" >> /etc/ssh/ssh_config \
+ && sed -i 's/#\(StrictModes \).*/\1no/g' /etc/ssh/sshd_config \
+ && rm -rf /var/lib/apt/lists/* \
+ && apt-get clean
+
 # install Python
 RUN wget -q https://www.python.org/ftp/python/$PYTHON_VERSION/Python-$PYTHON_VERSION.tgz \
  && tar -xzf Python-$PYTHON_VERSION.tgz \
@@ -122,76 +116,31 @@ ENV PATH="$PATH:/home/.openmpi/bin"
 ENV LD_LIBRARY_PATH="$LD_LIBRARY_PATH:/home/.openmpi/lib/"
 RUN ompi_info --parsable --all | grep mpi_built_with_cuda_support:value
 
-# Copy workaround script for incorrect hostname
-COPY changehostname.c /
-COPY start_with_right_hostname.sh /usr/local/bin/start_with_right_hostname.sh
+RUN mkdir -p /etc/pki/tls/certs && cp /etc/ssl/certs/ca-certificates.crt /etc/pki/tls/certs/ca-bundle.crt
 
-RUN ${PIP} install --no-cache-dir -U \
-    "bokeh>=2.3,<3" \
-    "awscli<2" \
-    scipy \
-    click \
-    "cryptography" \
-    "sagemaker>=2,<2.184" \
-    "sagemaker-pytorch-training" \
-    psutil==5.6.7 \
-    dataset \
-    transformers==4.36.2 \
-    Pillow
+# Install Neuron Driver, Runtime and Tools 
+RUN echo "deb https://apt.repos.neuron.amazonaws.com focal main" > /etc/apt/sources.list.d/neuron.list
+RUN wget -qO - https://apt.repos.neuron.amazonaws.com/GPG-PUB-KEY-AMAZON-AWS-NEURON.PUB | apt-key add -
 
-RUN mkdir -p /etc/pki/tls/certs && cp /etc/ssl/certs/ca-certificates.crt /etc/pki/tls/certs/ca-bundle.crt
+RUN apt-get update \
+ && apt-get install -y \
+    aws-neuronx-tools=$NEURONX_TOOLS_VERSION \
+    aws-neuronx-collectives=$NEURONX_COLLECTIVES_LIB_VERSION \
+    aws-neuronx-runtime-lib=$NEURONX_RUNTIME_LIB_VERSION \
+ && rm -rf /var/lib/apt/lists/* \
+ && rm -rf /tmp/tmp* \
+ && apt-get clean
+
+# Add Neuron PATH
+ENV PATH="/opt/aws/neuron/bin:${PATH}"
+
+# Install AWS CLI
+RUN ${PIP} install --no-cache-dir -U "awscli<2"
+
+# Install JAX & Neuron CC
 RUN ${PIP} config set global.extra-index-url https://pip.repos.neuron.amazonaws.com \
- && ${PIP} install --force-reinstall torch-neuronx==$NEURONX_FRAMEWORK_VERSION --extra-index-url https://pip.repos.neuron.amazonaws.com \
- && ${PIP} install --force-reinstall neuronx-cc==$NEURONX_CC_VERSION --extra-index-url https://pip.repos.neuron.amazonaws.com
-
-RUN ${PIP} install --force-reinstall --no-deps neuronx_distributed==$NEURONX_DISTRIBUTED_VERSION --extra-index-url https://pip.repos.neuron.amazonaws.com
-
-## Installation for Neuronx Distributed Training framework
-# Install Cython
-RUN pip install --no-cache-dir Cython
-
-# Copy the apex_setup.py file
-COPY apex_setup.py /root/apex_setup.py
-
-# Clone and build Apex
-RUN git clone https://github.com/NVIDIA/apex.git /root/apex \
-    && cd /root/apex \
-    && git checkout 23.05 \
-    && cp /root/apex_setup.py setup.py \
-    && python3 setup.py bdist_wheel
-
-#Install dependencies from requirements and extras for SageMaker usecase
-RUN wget https://raw.githubusercontent.com/aws-neuron/neuronx-distributed-training/master/requirements.txt \
-    && pip install --no-cache-dir -r requirements.txt /root/apex/dist/apex-0.1-py3-none-any.whl \
-    && pip install --force-reinstall "multiprocess==0.70.16" \
-    "dill==0.3.8" \
-    "torch==1.13.1"
-    
-
-RUN ${PIP} install --force-reinstall --no-deps neuronx_distributed_training==$NEURONX_DISTRIBUTED_TRAINING_VERSION --extra-index-url https://pip.repos.neuron.amazonaws.com
-
-# attrs, neuronx-cc required: >=19.2.0, sagemaker <24,>=23.1.0
-# protobuf neuronx-cc<4, sagemaker-training >=3.9.2,<=3.20.3
-# awscli 1.25.47 has requirement docutils<0.17,>=0.10
-# etcd for kubernetes installation
-# awscli 1.27.127 has requirement rsa<4.8,>=3.1.2, but you have rsa 4.9.
-# awscli 1.27.127 requires urllib3 < 1.27, python-etcd requires urllib3 >= 1.7, latest urllib3 release is 2.0.2
-RUN ${PIP} install --no-cache-dir -U \
-    "attrs<24,>=23.1.0" \
-    "protobuf>=3.18.3,<=3.20.3" \
-    "docutils>=0.10,<0.17" \
-    "rsa<4.8,>=3.1.2" \
-    "python-etcd" \
-    "urllib3>=1.26.0,<1.27"
-
-# Install extra packages needed by sagemaker (for passing test_utility_packages_using_import)
-RUN pip install --no-cache-dir -U \
-    "bokeh>=3.0.1,<4" \
-    "imageio>=2.22,<3" \
-    "opencv-python>=4.8.1.78" \
-    "plotly>=5.11,<6" \
-    "seaborn>=0.12,<1" \
-    "shap>=0.41,<1"
+ && ${PIP} install --force-reinstall neuronx-cc==$NEURONX_CC_VERSION --extra-index-url https://pip.repos.neuron.amazonaws.com \
+ && ${PIP} install --force-reinstall jax-neuronx==$NEURONX_JAX_TRAINING_VERSION --extra-index-url https://pip.repos.neuron.amazonaws.com
 
 # EFA Installer does apt get. Make sure to run apt update before that
 RUN apt-get update
@@ -205,27 +154,14 @@ RUN cd $HOME \
  && ./efa_installer.sh -y -g -d --skip-kmod --skip-limit-conf --no-verify \
  && cd $HOME
 
-
 # Clean up after apt update
 RUN rm -rf /var/lib/apt/lists/* \
  && rm -rf /tmp/tmp* \
  && apt-get clean
 
-# Install some common packages used by training scripts
-# torchvision needed for MLP. since it depends on torch and torch neuron/torch
-# is already installed install it with nodeps
-RUN pip3 install --no-cache-dir --no-deps -U \
-    torchvision==0.14.*
-
-# Needed for running bert training scripts
-RUN pip3 install --no-cache-dir -U \
-    graphviz \
-    tensorboard==2.6 \
-    accelerate \
-    sentencepiece!=0.1.92 \
-    h5py \
-    requests
-
+# Copy workaround script for incorrect hostname
+COPY changehostname.c /
+COPY start_with_right_hostname.sh /usr/local/bin/start_with_right_hostname.sh
 COPY deep_learning_container.py /usr/local/bin/deep_learning_container.py
 
 RUN chmod +x /usr/local/bin/start_with_right_hostname.sh \
@@ -241,8 +177,8 @@ RUN HOME_DIR=/root \
  && rm -rf ${HOME_DIR}/oss_compliance* \
  && rm -rf /tmp/tmp*
 
-RUN curl -o /license.txt  https://aws-dlc-licenses.s3.amazonaws.com/pytorch-1.13/license.txt
-
 # Starts framework
 ENTRYPOINT ["bash", "-m", "start_with_right_hostname.sh"]
 CMD ["/bin/bash"]
+
+HEALTHCHECK CMD curl --fail http://localhost:8080/ping || exit 1
\ No newline at end of file
diff --git a/docker/jax/training/0.4/Dockerfile.neuronx.cve_allowlist.json b/docker/jax/training/0.4/Dockerfile.neuronx.cve_allowlist.json
new file mode 100644
index 0000000..a83234f
--- /dev/null
+++ b/docker/jax/training/0.4/Dockerfile.neuronx.cve_allowlist.json
@@ -0,0 +1,27 @@
+{
+    "CVE-2024-35195": {
+        "description": "Requests is a HTTP library. Prior to 2.32.0, when making requests through a Requests `Session`, if the first request is made with `verify=False` to disable cert verification, all subsequent requests to the same host will continue to ignore cert verification regardless of changes to the value of `verify`. This behavior will continue for the lifecycle of the connection in the connection pool. This vulnerability is fixed in 2.32.0.",
+        "remediation": {
+            "recommendation": {
+                "text": "None Provided"
+            }
+        },
+        "score": 0.0,
+        "score_details": {},
+        "severity": "UNTRIAGED",
+        "source": "NVD",
+        "source_url": "https://nvd.nist.gov/vuln/detail/CVE-2024-35195",
+        "status": "ACTIVE",
+        "title": "CVE-2024-35195 - requests",
+        "vulnerability_id": "CVE-2024-35195",
+        "vulnerable_packages": [
+            {
+                "epoch": 0,
+                "filePath": "usr/local/lib/python3.10/site-packages/requests-2.31.0.dist-info/METADATA",
+                "name": "requests",
+                "packageManager": "PYTHONPKG",
+                "version": "2.31.0"
+            }
+        ]
+    }
+}
\ No newline at end of file
diff --git a/docker/pytorch/inference/1.13.1/Dockerfile.neuron b/docker/pytorch/inference/1.13.1/Dockerfile.neuron
deleted file mode 100644
index 7fddb78..0000000
--- a/docker/pytorch/inference/1.13.1/Dockerfile.neuron
+++ /dev/null
@@ -1,153 +0,0 @@
-FROM public.ecr.aws/docker/library/ubuntu:20.04
-
-LABEL dlc_major_version="1"
-LABEL maintainer="Amazon AI"
-LABEL com.amazonaws.sagemaker.capabilities.accept-bind-to-port=true
-
-# Neuron SDK components version numbers
-ARG NEURON_FRAMEWORK_VERSION=1.13.1.2.11.13.0
-ARG NEURON_CC_VERSION=1.24.0.0
-ARG NEURONX_TOOLS_VERSION=2.19.0.0
-
-ARG PYTHON=python3.10
-ARG PYTHON_VERSION=3.10.12
-ARG TORCHSERVE_VERSION=0.11.0
-ARG SM_TOOLKIT_VERSION=2.0.21
-ARG MAMBA_VERSION=23.1.0-4
-
-# See http://bugs.python.org/issue19846
-ENV LANG=C.UTF-8
-ENV LD_LIBRARY_PATH=/lib/x86_64-linux-gnu:/opt/conda/lib/:$LD_LIBRARY_PATH
-ENV PATH=/opt/conda/bin:/opt/aws/neuron/bin:$PATH
-ENV SAGEMAKER_SERVING_MODULE=sagemaker_pytorch_serving_container.serving:main
-ENV TEMP=/home/model-server/tmp
-
-RUN apt-get update \
- && apt-get upgrade -y \
- && apt-get install -y --no-install-recommends software-properties-common \
- && add-apt-repository ppa:openjdk-r/ppa \
- && apt-get update \
- && apt-get install -y --no-install-recommends \
-    build-essential \
-    apt-transport-https \
-    ca-certificates \
-    cmake \
-    curl \
-    emacs \
-    git \
-    jq \
-    libgl1-mesa-glx \
-    libglib2.0-0 \
-    libsm6 \
-    libxext6 \
-    libxrender-dev \
-    openjdk-11-jdk \
-    vim \
-    wget \
-    unzip \
-    zlib1g-dev \
-    libcap-dev \
-    gnupg2 \
-    gpg-agent \
- && rm -rf /var/lib/apt/lists/* \
- && rm -rf /tmp/tmp* \
- && apt-get clean
-
-RUN echo "deb https://apt.repos.neuron.amazonaws.com focal main" > /etc/apt/sources.list.d/neuron.list
-RUN wget -qO - https://apt.repos.neuron.amazonaws.com/GPG-PUB-KEY-AMAZON-AWS-NEURON.PUB | apt-key add -
-
-RUN apt-get update \
- && apt-get install -y aws-neuronx-tools=$NEURONX_TOOLS_VERSION \
- && rm -rf /var/lib/apt/lists/* \
- && rm -rf /tmp/tmp* \
- && apt-get clean
-
-# https://github.com/docker-library/openjdk/issues/261 https://github.com/docker-library/openjdk/pull/263/files
-RUN keytool -importkeystore -srckeystore /etc/ssl/certs/java/cacerts -destkeystore /etc/ssl/certs/java/cacerts.jks -deststoretype JKS -srcstorepass changeit -deststorepass changeit -noprompt; \
-    mv /etc/ssl/certs/java/cacerts.jks /etc/ssl/certs/java/cacerts; \
-    /var/lib/dpkg/info/ca-certificates-java.postinst configure;
-
-RUN curl -L -o ~/mambaforge.sh https://github.com/conda-forge/miniforge/releases/download/${MAMBA_VERSION}/Mambaforge-${MAMBA_VERSION}-Linux-x86_64.sh \
- && chmod +x ~/mambaforge.sh \
- && ~/mambaforge.sh -b -p /opt/conda \
- && rm ~/mambaforge.sh \
- && /opt/conda/bin/conda update -y conda \
- && /opt/conda/bin/conda install -c conda-forge -y \
-    python=$PYTHON_VERSION \
-    pyopenssl \
-    cython \
-    mkl-include \
-    mkl \
-    parso \
-    typing \
-    # Below 2 are included in miniconda base, but not mamba so need to install
-    conda-content-trust \
-    charset-normalizer \
-
- && /opt/conda/bin/conda clean -ya
-
-RUN conda install -c conda-forge \
-    scikit-learn \
-    h5py \
-    requests \
- && conda clean -ya \
- && pip install --upgrade pip --trusted-host pypi.org --trusted-host files.pythonhosted.org \
- && ln -s /opt/conda/bin/pip /usr/local/bin/pip3 \
- && pip install packaging \
-    enum-compat \
-    ipython
-
-RUN pip install --no-cache-dir -U \
-    opencv-python>=4.8.1.78 \
-    "numpy<1.24,>1.21" \
-    "scipy>=1.8.0" \
-    six \
-    "pillow>=10.0.1" \
-    "awscli<2" \
-    pandas==1.* \
-    boto3 \
-    cryptography
-
-RUN pip install neuron-cc==$NEURON_CC_VERSION --extra-index-url https://pip.repos.neuron.amazonaws.com \
-    torch-neuron==$NEURON_FRAMEWORK_VERSION \
- && pip install -U protobuf==3.19.5 \
-    torchserve==${TORCHSERVE_VERSION} \
-    torch-model-archiver==${TORCHSERVE_VERSION} \
- && pip install --no-deps --no-cache-dir -U torchvision==0.14.*
-
-RUN useradd -m model-server \
- && mkdir -p /home/model-server/tmp /opt/ml/model \
- && chown -R model-server /home/model-server /opt/ml/model
-
-COPY neuron-entrypoint.py /usr/local/bin/dockerd-entrypoint.py
-COPY neuron-monitor.sh /usr/local/bin/neuron-monitor.sh
-COPY torchserve-neuron.sh /usr/local/bin/entrypoint.sh
-COPY config.properties /home/model-server
-
-RUN chmod +x /usr/local/bin/dockerd-entrypoint.py \
- && chmod +x /usr/local/bin/neuron-monitor.sh \
- && chmod +x /usr/local/bin/entrypoint.sh
-
-ADD https://raw.githubusercontent.com/aws/deep-learning-containers/master/src/deep_learning_container.py /usr/local/bin/deep_learning_container.py
-
-RUN chmod +x /usr/local/bin/deep_learning_container.py
-
-RUN pip install --no-cache-dir "sagemaker-pytorch-inference==${SM_TOOLKIT_VERSION}"
-
-RUN HOME_DIR=/root \
- && curl -o ${HOME_DIR}/oss_compliance.zip https://aws-dlinfra-utilities.s3.amazonaws.com/oss_compliance.zip \
- && unzip ${HOME_DIR}/oss_compliance.zip -d ${HOME_DIR}/ \
- && cp ${HOME_DIR}/oss_compliance/test/testOSSCompliance /usr/local/bin/testOSSCompliance \
- && chmod +x /usr/local/bin/testOSSCompliance \
- && chmod +x ${HOME_DIR}/oss_compliance/generate_oss_compliance.sh \
- && ${HOME_DIR}/oss_compliance/generate_oss_compliance.sh ${HOME_DIR} ${PYTHON} \
- && rm -rf ${HOME_DIR}/oss_compliance* \
- # conda leaves an empty /root/.cache/conda/notices.cache file which is not removed by conda clean -ya
- && rm -rf ${HOME_DIR}/.cache/conda
-
-RUN curl https://aws-dlc-licenses.s3.amazonaws.com/pytorch-1.13/license.txt -o /license.txt
-
-EXPOSE 8080 8081
-
-ENTRYPOINT ["python", "/usr/local/bin/dockerd-entrypoint.py"]
-CMD ["/usr/local/bin/entrypoint.sh"]
diff --git a/docker/pytorch/inference/1.13.1/Dockerfile.neuron.cve_allowlist.json b/docker/pytorch/inference/1.13.1/Dockerfile.neuron.cve_allowlist.json
deleted file mode 100644
index f9abeb3..0000000
--- a/docker/pytorch/inference/1.13.1/Dockerfile.neuron.cve_allowlist.json
+++ /dev/null
@@ -1,127 +0,0 @@
-{
-    "CVE-2024-2511": {
-        "description": "Issue summary: Some non-default TLS server configurations can cause unbounded memory growth when processing TLSv1.3 sessions\n\nImpact summary: An attacker may exploit certain server configurations to trigger unbounded memory growth that would lead to a Denial of Service\n\nThis problem can occur in TLSv1.3 if the non-default SSL_OP_NO_TICKET option is being used (but not if early_data support is also configured and the default anti-replay protection is in use). In this case, under certain conditions, the session cache can get into an incorrect state and it will fail to flush properly as it fills. The session cache will continue to grow in an unbounded manner. A malicious client could deliberately create the scenario for this failure to force a Denial of Service. It may also happen by accident in normal operation.\n\nThis issue only affects TLS servers supporting TLSv1.3. It does not affect TLS clients.\n\nThe FIPS modules in 3.2, 3.1 and 3.0 are not affected by this issue. OpenSSL 1.0.2 is also not affected by this ",
-        "remediation": {
-            "recommendation": {
-                "text": "None Provided"
-            }
-        },
-        "score": 0.0,
-        "score_details": {},
-        "severity": "UNTRIAGED",
-        "source": "NVD",
-        "source_url": "https://nvd.nist.gov/vuln/detail/CVE-2024-2511",
-        "status": "ACTIVE",
-        "title": "CVE-2024-2511 - pyOpenSSL",
-        "vulnerability_id": "CVE-2024-2511",
-        "vulnerable_packages": [
-            {
-                "epoch": 0,
-                "filePath": "opt/conda/lib/python3.10/site-packages/pyOpenSSL-24.2.1.dist-info/METADATA",
-                "name": "pyOpenSSL",
-                "packageManager": "PYTHONPKG",
-                "version": "24.2.1"
-            }
-        ]
-    },
-    "CVE-2024-31580": {
-        "description": "PyTorch before v2.2.0 was discovered to contain a heap buffer overflow vulnerability in the component /runtime/vararg_functions.cpp. This vulnerability allows attackers to cause a Denial of Service (DoS) via a crafted input.",
-        "remediation": {
-            "recommendation": {
-                "text": "None Provided"
-            }
-        },
-        "score": 0.0,
-        "score_details": {},
-        "severity": "UNTRIAGED",
-        "source": "NVD",
-        "source_url": "https://nvd.nist.gov/vuln/detail/CVE-2024-31580",
-        "status": "ACTIVE",
-        "title": "CVE-2024-31580 - torch",
-        "vulnerability_id": "CVE-2024-31580",
-        "vulnerable_packages": [
-            {
-                "epoch": 0,
-                "filePath": "opt/conda/lib/python3.10/site-packages/torch-1.13.1.dist-info/METADATA",
-                "name": "torch",
-                "packageManager": "PYTHONPKG",
-                "version": "1.13.1"
-            }
-        ]
-    },
-    "CVE-2024-31583": {
-        "description": "Pytorch before version v2.2.0 was discovered to contain a use-after-free vulnerability in torch/csrc/jit/mobile/interpreter.cpp.",
-        "remediation": {
-            "recommendation": {
-                "text": "None Provided"
-            }
-        },
-        "score": 0.0,
-        "score_details": {},
-        "severity": "UNTRIAGED",
-        "source": "NVD",
-        "source_url": "https://nvd.nist.gov/vuln/detail/CVE-2024-31583",
-        "status": "ACTIVE",
-        "title": "CVE-2024-31583 - torch",
-        "vulnerability_id": "CVE-2024-31583",
-        "vulnerable_packages": [
-            {
-                "epoch": 0,
-                "filePath": "opt/conda/lib/python3.10/site-packages/torch-1.13.1.dist-info/METADATA",
-                "name": "torch",
-                "packageManager": "PYTHONPKG",
-                "version": "1.13.1"
-            }
-        ]
-    },
-    "CVE-2024-37891": {
-        "description": "urllib3 is a user-friendly HTTP client library for Python. When using urllib3's proxy support with `ProxyManager`, the `Proxy-Authorization` header is only sent to the configured proxy, as expected. However, when sending HTTP requests *without* using urllib3's proxy support, it's possible to accidentally configure the `Proxy-Authorization` header even though it won't have any effect as the request is not using a forwarding proxy or a tunneling proxy. In those cases, urllib3 doesn't treat the `Proxy-Authorization` HTTP header as one carrying authentication material and thus doesn't strip the header on cross-origin redirects. Because this is a highly unlikely scenario, we believe the severity of this vulnerability is low for almost all users. Out of an abundance of caution urllib3 will automatically strip the `Proxy-Authorization` header during cross-origin redirects to avoid the small chance that users are doing this on accident. Users should use urllib3's proxy support or disable automatic redirects to achiev",
-        "remediation": {
-            "recommendation": {
-                "text": "None Provided"
-            }
-        },
-        "score": 0.0,
-        "score_details": {},
-        "severity": "UNTRIAGED",
-        "source": "NVD",
-        "source_url": "https://nvd.nist.gov/vuln/detail/CVE-2024-37891",
-        "status": "ACTIVE",
-        "title": "CVE-2024-37891 - urllib3",
-        "vulnerability_id": "CVE-2024-37891",
-        "vulnerable_packages": [
-            {
-                "epoch": 0,
-                "filePath": "opt/conda/lib/python3.10/site-packages/urllib3-2.0.7.dist-info/METADATA",
-                "name": "urllib3",
-                "packageManager": "PYTHONPKG",
-                "version": "2.0.7"
-            }
-        ]
-    },
-    "CVE-2024-6345": {
-        "description": "A vulnerability in the package_index module of pypa/setuptools versions up to 69.1.1 allows for remote code execution via its download functions. These functions, which are used to download packages from URLs provided by users or retrieved from package index servers, are susceptible to code injection. If these functions are exposed to user-controlled inputs, such as package URLs, they can execute arbitrary commands on the system. The issue is fixed in version 70.0.",
-        "remediation": {
-            "recommendation": {
-                "text": "None Provided"
-            }
-        },
-        "score": 0.0,
-        "score_details": {},
-        "severity": "UNTRIAGED",
-        "source": "NVD",
-        "source_url": "https://nvd.nist.gov/vuln/detail/CVE-2024-6345",
-        "status": "ACTIVE",
-        "title": "CVE-2024-6345 - setuptools",
-        "vulnerability_id": "CVE-2024-6345",
-        "vulnerable_packages": [
-            {
-                "epoch": 0,
-                "filePath": "opt/conda/lib/python3.10/site-packages/setuptools-69.5.1.dist-info/METADATA",
-                "name": "setuptools",
-                "packageManager": "PYTHONPKG",
-                "version": "69.5.1"
-            }
-        ]
-    }
-}
\ No newline at end of file
diff --git a/docker/pytorch/inference/1.13.1/Dockerfile.neuronx b/docker/pytorch/inference/1.13.1/Dockerfile.neuronx
deleted file mode 100644
index d4b9389..0000000
--- a/docker/pytorch/inference/1.13.1/Dockerfile.neuronx
+++ /dev/null
@@ -1,168 +0,0 @@
-FROM public.ecr.aws/docker/library/ubuntu:20.04
-
-LABEL dlc_major_version="1"
-LABEL maintainer="Amazon AI"
-LABEL com.amazonaws.sagemaker.capabilities.accept-bind-to-port=true
-
-# Neuron SDK components version numbers
-ARG NEURONX_FRAMEWORK_VERSION=1.13.1.1.16.0
-ARG NEURONX_DISTRIBUTED_VERSION=0.9.0
-ARG NEURONX_CC_VERSION=2.15.143.0
-ARG NEURONX_TRANSFORMERS_VERSION=0.12.313
-ARG NEURONX_COLLECTIVES_LIB_VERSION=2.22.33.0-d2128d1aa
-ARG NEURONX_RUNTIME_LIB_VERSION=2.22.19.0-5856c0b42
-ARG NEURONX_TOOLS_VERSION=2.19.0.0
-
-ARG PYTHON=python3.10
-ARG PYTHON_VERSION=3.10.12
-ARG TORCHSERVE_VERSION=0.11.0
-ARG SM_TOOLKIT_VERSION=2.0.21
-ARG MAMBA_VERSION=23.1.0-4
-
-# See http://bugs.python.org/issue19846
-ENV LANG=C.UTF-8
-ENV LD_LIBRARY_PATH=/opt/aws/neuron/lib:/lib/x86_64-linux-gnu:/opt/conda/lib/:$LD_LIBRARY_PATH
-ENV PATH=/opt/conda/bin:/opt/aws/neuron/bin:$PATH
-ENV SAGEMAKER_SERVING_MODULE=sagemaker_pytorch_serving_container.serving:main
-ENV TEMP=/home/model-server/tmp
-
-RUN apt-get update \
- && apt-get upgrade -y \
- && apt-get install -y --no-install-recommends software-properties-common \
- && add-apt-repository ppa:openjdk-r/ppa \
- && apt-get update \
- && apt-get install -y --no-install-recommends \
-    build-essential \
-    apt-transport-https \
-    ca-certificates \
-    cmake \
-    curl \
-    emacs \
-    git \
-    jq \
-    libgl1-mesa-glx \
-    libglib2.0-0 \
-    libsm6 \
-    libxext6 \
-    libxrender-dev \
-    openjdk-11-jdk \
-    vim \
-    wget \
-    unzip \
-    zlib1g-dev \
-    libcap-dev \
-    gnupg2 \
-    gpg-agent \
- && rm -rf /var/lib/apt/lists/* \
- && rm -rf /tmp/tmp* \
- && apt-get clean
-
-RUN echo "deb https://apt.repos.neuron.amazonaws.com focal main" > /etc/apt/sources.list.d/neuron.list
-RUN wget -qO - https://apt.repos.neuron.amazonaws.com/GPG-PUB-KEY-AMAZON-AWS-NEURON.PUB | apt-key add -
-
-RUN apt-get update \
- && apt-get install -y \
-    aws-neuronx-tools=$NEURONX_TOOLS_VERSION \
-    aws-neuronx-collectives=$NEURONX_COLLECTIVES_LIB_VERSION \
-    aws-neuronx-runtime-lib=$NEURONX_RUNTIME_LIB_VERSION \
- && rm -rf /var/lib/apt/lists/* \
- && rm -rf /tmp/tmp* \
- && apt-get clean
-
-# https://github.com/docker-library/openjdk/issues/261 https://github.com/docker-library/openjdk/pull/263/files
-RUN keytool -importkeystore -srckeystore /etc/ssl/certs/java/cacerts -destkeystore /etc/ssl/certs/java/cacerts.jks -deststoretype JKS -srcstorepass changeit -deststorepass changeit -noprompt; \
-    mv /etc/ssl/certs/java/cacerts.jks /etc/ssl/certs/java/cacerts; \
-    /var/lib/dpkg/info/ca-certificates-java.postinst configure;
-
-RUN curl -L -o ~/mambaforge.sh https://github.com/conda-forge/miniforge/releases/download/${MAMBA_VERSION}/Mambaforge-${MAMBA_VERSION}-Linux-x86_64.sh \
- && chmod +x ~/mambaforge.sh \
- && ~/mambaforge.sh -b -p /opt/conda \
- && rm ~/mambaforge.sh \
- && /opt/conda/bin/conda update -y conda \
- && /opt/conda/bin/conda install -c conda-forge -y \
-    python=$PYTHON_VERSION \
-    pyopenssl \
-    cython \
-    mkl-include \
-    mkl \
-    parso \
-    typing \
-    # Below 2 are included in miniconda base, but not mamba so need to install
-    conda-content-trust \
-    charset-normalizer \
-
- && /opt/conda/bin/conda clean -ya
-
-RUN conda install -c conda-forge \
-    scikit-learn \
-    h5py \
-    requests \
- && conda clean -ya \
- && pip install --upgrade pip --trusted-host pypi.org --trusted-host files.pythonhosted.org \
- && ln -s /opt/conda/bin/pip /usr/local/bin/pip3 \
- && pip install packaging \
-    enum-compat \
-    ipython
-
-RUN pip install --no-cache-dir -U \
-    opencv-python>=4.8.1.78 \
-    "numpy<1.24,>1.21" \
-    "scipy>=1.8.0" \
-    six \
-    "pillow>=10.0.1" \
-    "awscli<2" \
-    pandas==1.* \
-    boto3 \
-    cryptography
-
-RUN pip install -U --extra-index-url https://pip.repos.neuron.amazonaws.com \
-    neuronx-cc==$NEURONX_CC_VERSION \
-    torch-neuronx==$NEURONX_FRAMEWORK_VERSION \
-    transformers-neuronx==$NEURONX_TRANSFORMERS_VERSION \
- && pip install -U "protobuf>=3.18.3,<4" \
-    torchserve==${TORCHSERVE_VERSION} \
-    torch-model-archiver==${TORCHSERVE_VERSION} \
- && pip install --no-deps --no-cache-dir -U torchvision==0.14.* \
- && pip install --no-deps -U --extra-index-url https://pip.repos.neuron.amazonaws.com neuronx_distributed==$NEURONX_DISTRIBUTED_VERSION
-
-RUN useradd -m model-server \
- && mkdir -p /home/model-server/tmp /opt/ml/model \
- && chown -R model-server /home/model-server /opt/ml/model
-
-COPY neuron-entrypoint.py /usr/local/bin/dockerd-entrypoint.py
-COPY neuron-monitor.sh /usr/local/bin/neuron-monitor.sh
-COPY torchserve-neuron.sh /usr/local/bin/entrypoint.sh
-COPY config.properties /home/model-server
-
-RUN chmod +x /usr/local/bin/dockerd-entrypoint.py \
- && chmod +x /usr/local/bin/neuron-monitor.sh \
- && chmod +x /usr/local/bin/entrypoint.sh
-
-ADD https://raw.githubusercontent.com/aws/deep-learning-containers/master/src/deep_learning_container.py /usr/local/bin/deep_learning_container.py
-
-RUN chmod +x /usr/local/bin/deep_learning_container.py
-
-RUN pip install --no-cache-dir "sagemaker-pytorch-inference==${SM_TOOLKIT_VERSION}"
-
-# patch default_pytorch_inference_handler.py to import torch_neuronx
-RUN DEST_DIR=$(python -c "import os.path, sagemaker_pytorch_serving_container; print(os.path.dirname(sagemaker_pytorch_serving_container.__file__))") \
- && DEST_FILE=${DEST_DIR}/default_pytorch_inference_handler.py \
- && sed -i "s/import torch/import torch, torch_neuronx/" ${DEST_FILE}
-
-RUN HOME_DIR=/root \
- && curl -o ${HOME_DIR}/oss_compliance.zip https://aws-dlinfra-utilities.s3.amazonaws.com/oss_compliance.zip \
- && unzip ${HOME_DIR}/oss_compliance.zip -d ${HOME_DIR}/ \
- && cp ${HOME_DIR}/oss_compliance/test/testOSSCompliance /usr/local/bin/testOSSCompliance \
- && chmod +x /usr/local/bin/testOSSCompliance \
- && chmod +x ${HOME_DIR}/oss_compliance/generate_oss_compliance.sh \
- && ${HOME_DIR}/oss_compliance/generate_oss_compliance.sh ${HOME_DIR} ${PYTHON} \
- && rm -rf ${HOME_DIR}/oss_compliance* \
- # conda leaves an empty /root/.cache/conda/notices.cache file which is not removed by conda clean -ya
- && rm -rf ${HOME_DIR}/.cache/conda
-
-RUN curl https://aws-dlc-licenses.s3.amazonaws.com/pytorch-1.13/license.txt -o /license.txt
-
-EXPOSE 8080 8081
-
-ENTRYPOINT ["python", "/usr/local/bin/dockerd-entrypoint.py"]
-CMD ["/usr/local/bin/entrypoint.sh"]
diff --git a/docker/pytorch/inference/1.13.1/Dockerfile.neuronx.cve_allowlist.json b/docker/pytorch/inference/1.13.1/Dockerfile.neuronx.cve_allowlist.json
deleted file mode 100644
index 0d9dfa2..0000000
--- a/docker/pytorch/inference/1.13.1/Dockerfile.neuronx.cve_allowlist.json
+++ /dev/null
@@ -1,152 +0,0 @@
-{
-    "CVE-2024-2511": {
-        "description": "Issue summary: Some non-default TLS server configurations can cause unbounded memory growth when processing TLSv1.3 sessions\n\nImpact summary: An attacker may exploit certain server configurations to trigger unbounded memory growth that would lead to a Denial of Service\n\nThis problem can occur in TLSv1.3 if the non-default SSL_OP_NO_TICKET option is being used (but not if early_data support is also configured and the default anti-replay protection is in use). In this case, under certain conditions, the session cache can get into an incorrect state and it will fail to flush properly as it fills. The session cache will continue to grow in an unbounded manner. A malicious client could deliberately create the scenario for this failure to force a Denial of Service. It may also happen by accident in normal operation.\n\nThis issue only affects TLS servers supporting TLSv1.3. It does not affect TLS clients.\n\nThe FIPS modules in 3.2, 3.1 and 3.0 are not affected by this issue. OpenSSL 1.0.2 is also not affected by this ",
-        "remediation": {
-            "recommendation": {
-                "text": "None Provided"
-            }
-        },
-        "score": 0.0,
-        "score_details": {},
-        "severity": "UNTRIAGED",
-        "source": "NVD",
-        "source_url": "https://nvd.nist.gov/vuln/detail/CVE-2024-2511",
-        "status": "ACTIVE",
-        "title": "CVE-2024-2511 - pyOpenSSL",
-        "vulnerability_id": "CVE-2024-2511",
-        "vulnerable_packages": [
-            {
-                "epoch": 0,
-                "filePath": "opt/conda/lib/python3.10/site-packages/pyOpenSSL-24.2.1.dist-info/METADATA",
-                "name": "pyOpenSSL",
-                "packageManager": "PYTHONPKG",
-                "version": "24.2.1"
-            }
-        ]
-    },
-    "CVE-2024-31580": {
-        "description": "PyTorch before v2.2.0 was discovered to contain a heap buffer overflow vulnerability in the component /runtime/vararg_functions.cpp. This vulnerability allows attackers to cause a Denial of Service (DoS) via a crafted input.",
-        "remediation": {
-            "recommendation": {
-                "text": "None Provided"
-            }
-        },
-        "score": 0.0,
-        "score_details": {},
-        "severity": "UNTRIAGED",
-        "source": "NVD",
-        "source_url": "https://nvd.nist.gov/vuln/detail/CVE-2024-31580",
-        "status": "ACTIVE",
-        "title": "CVE-2024-31580 - torch",
-        "vulnerability_id": "CVE-2024-31580",
-        "vulnerable_packages": [
-            {
-                "epoch": 0,
-                "filePath": "opt/conda/lib/python3.10/site-packages/torch-1.13.1.dist-info/METADATA",
-                "name": "torch",
-                "packageManager": "PYTHONPKG",
-                "version": "1.13.1"
-            }
-        ]
-    },
-    "CVE-2024-31583": {
-        "description": "Pytorch before version v2.2.0 was discovered to contain a use-after-free vulnerability in torch/csrc/jit/mobile/interpreter.cpp.",
-        "remediation": {
-            "recommendation": {
-                "text": "None Provided"
-            }
-        },
-        "score": 0.0,
-        "score_details": {},
-        "severity": "UNTRIAGED",
-        "source": "NVD",
-        "source_url": "https://nvd.nist.gov/vuln/detail/CVE-2024-31583",
-        "status": "ACTIVE",
-        "title": "CVE-2024-31583 - torch",
-        "vulnerability_id": "CVE-2024-31583",
-        "vulnerable_packages": [
-            {
-                "epoch": 0,
-                "filePath": "opt/conda/lib/python3.10/site-packages/torch-1.13.1.dist-info/METADATA",
-                "name": "torch",
-                "packageManager": "PYTHONPKG",
-                "version": "1.13.1"
-            }
-        ]
-    },
-    "CVE-2024-35195": {
-        "description": "Requests is a HTTP library. Prior to 2.32.0, when making requests through a Requests `Session`, if the first request is made with `verify=False` to disable cert verification, all subsequent requests to the same host will continue to ignore cert verification regardless of changes to the value of `verify`. This behavior will continue for the lifecycle of the connection in the connection pool. This vulnerability is fixed in 2.32.0.",
-        "remediation": {
-            "recommendation": {
-                "text": "None Provided"
-            }
-        },
-        "score": 0.0,
-        "score_details": {},
-        "severity": "UNTRIAGED",
-        "source": "NVD",
-        "source_url": "https://nvd.nist.gov/vuln/detail/CVE-2024-35195",
-        "status": "ACTIVE",
-        "title": "CVE-2024-35195 - requests",
-        "vulnerability_id": "CVE-2024-35195",
-        "vulnerable_packages": [
-            {
-                "epoch": 0,
-                "filePath": "opt/conda/lib/python3.10/site-packages/requests-2.31.0.dist-info/METADATA",
-                "name": "requests",
-                "packageManager": "PYTHONPKG",
-                "version": "2.31.0"
-            }
-        ]
-    },
-    "CVE-2024-37891": {
-        "description": "urllib3 is a user-friendly HTTP client library for Python. When using urllib3's proxy support with `ProxyManager`, the `Proxy-Authorization` header is only sent to the configured proxy, as expected. However, when sending HTTP requests *without* using urllib3's proxy support, it's possible to accidentally configure the `Proxy-Authorization` header even though it won't have any effect as the request is not using a forwarding proxy or a tunneling proxy. In those cases, urllib3 doesn't treat the `Proxy-Authorization` HTTP header as one carrying authentication material and thus doesn't strip the header on cross-origin redirects. Because this is a highly unlikely scenario, we believe the severity of this vulnerability is low for almost all users. Out of an abundance of caution urllib3 will automatically strip the `Proxy-Authorization` header during cross-origin redirects to avoid the small chance that users are doing this on accident. Users should use urllib3's proxy support or disable automatic redirects to achiev",
-        "remediation": {
-            "recommendation": {
-                "text": "None Provided"
-            }
-        },
-        "score": 0.0,
-        "score_details": {},
-        "severity": "UNTRIAGED",
-        "source": "NVD",
-        "source_url": "https://nvd.nist.gov/vuln/detail/CVE-2024-37891",
-        "status": "ACTIVE",
-        "title": "CVE-2024-37891 - urllib3",
-        "vulnerability_id": "CVE-2024-37891",
-        "vulnerable_packages": [
-            {
-                "epoch": 0,
-                "filePath": "opt/conda/lib/python3.10/site-packages/urllib3-2.0.7.dist-info/METADATA",
-                "name": "urllib3",
-                "packageManager": "PYTHONPKG",
-                "version": "2.0.7"
-            }
-        ]
-    },
-    "CVE-2024-6345": {
-        "description": "A vulnerability in the package_index module of pypa/setuptools versions up to 69.1.1 allows for remote code execution via its download functions. These functions, which are used to download packages from URLs provided by users or retrieved from package index servers, are susceptible to code injection. If these functions are exposed to user-controlled inputs, such as package URLs, they can execute arbitrary commands on the system. The issue is fixed in version 70.0.",
-        "remediation": {
-            "recommendation": {
-                "text": "None Provided"
-            }
-        },
-        "score": 0.0,
-        "score_details": {},
-        "severity": "UNTRIAGED",
-        "source": "NVD",
-        "source_url": "https://nvd.nist.gov/vuln/detail/CVE-2024-6345",
-        "status": "ACTIVE",
-        "title": "CVE-2024-6345 - setuptools",
-        "vulnerability_id": "CVE-2024-6345",
-        "vulnerable_packages": [
-            {
-                "epoch": 0,
-                "filePath": "opt/conda/lib/python3.10/site-packages/setuptools-69.5.1.dist-info/METADATA",
-                "name": "setuptools",
-                "packageManager": "PYTHONPKG",
-                "version": "69.5.1"
-            }
-        ]
-    }
-}
\ No newline at end of file
diff --git a/docker/pytorch/inference/2.1.2/Dockerfile.neuronx b/docker/pytorch/inference/2.5.1/Dockerfile.neuronx
similarity index 88%
rename from docker/pytorch/inference/2.1.2/Dockerfile.neuronx
rename to docker/pytorch/inference/2.5.1/Dockerfile.neuronx
index c8a8d0d..8d35e3d 100644
--- a/docker/pytorch/inference/2.1.2/Dockerfile.neuronx
+++ b/docker/pytorch/inference/2.5.1/Dockerfile.neuronx
@@ -1,17 +1,18 @@
-FROM public.ecr.aws/docker/library/ubuntu:20.04
+FROM public.ecr.aws/docker/library/ubuntu:22.04
 
 LABEL dlc_major_version="1"
 LABEL maintainer="Amazon AI"
 LABEL com.amazonaws.sagemaker.capabilities.accept-bind-to-port=true
 
 # Neuron SDK components version numbers
-ARG NEURONX_DISTRIBUTED_VERSION=0.9.0
-ARG NEURONX_CC_VERSION=2.15.143.0
-ARG NEURONX_FRAMEWORK_VERSION=2.1.2.2.3.2
-ARG NEURONX_TRANSFORMERS_VERSION=0.12.313
-ARG NEURONX_COLLECTIVES_LIB_VERSION=2.22.33.0-d2128d1aa
-ARG NEURONX_RUNTIME_LIB_VERSION=2.22.19.0-5856c0b42
-ARG NEURONX_TOOLS_VERSION=2.19.0.0
+ARG NEURONX_CC_VERSION=2.16.345.0
+ARG NEURONX_FRAMEWORK_VERSION=2.5.1.2.4.0
+ARG NEURONX_TRANSFORMERS_VERSION=0.13.322
+ARG NEURONX_COLLECTIVES_LIB_VERSION=2.23.133.0-3e70920f2
+ARG NEURONX_RUNTIME_LIB_VERSION=2.23.110.0-9b5179492
+ARG NEURONX_TOOLS_VERSION=2.20.204.0
+ARG NEURONX_DISTRIBUTED_VERSION=0.10.0
+ARG NEURONX_DISTRIBUTED_INFERENCE_VERSION=0.1.0
 
 ARG PYTHON=python3.10
 ARG PYTHON_VERSION=3.10.12
@@ -28,31 +29,29 @@ ENV TEMP=/home/model-server/tmp
 
 RUN apt-get update \
  && apt-get upgrade -y \
- && apt-get install -y --no-install-recommends software-properties-common \
- && add-apt-repository ppa:openjdk-r/ppa \
- && apt-get update \
  && apt-get install -y --no-install-recommends \
-    build-essential \
     apt-transport-https \
+    build-essential \
     ca-certificates \
     cmake \
     curl \
     emacs \
     git \
+    gnupg2 \
+    gpg-agent \
     jq \
     libgl1-mesa-glx \
     libglib2.0-0 \
     libsm6 \
     libxext6 \
     libxrender-dev \
+    libcap-dev \
+    libhwloc-dev \
     openjdk-11-jdk \
+    unzip \
     vim \
     wget \
-    unzip \
     zlib1g-dev \
-    libcap-dev \
-    gnupg2 \
-    gpg-agent \
  && rm -rf /var/lib/apt/lists/* \
  && rm -rf /tmp/tmp* \
  && apt-get clean
@@ -90,7 +89,6 @@ RUN curl -L -o ~/mambaforge.sh https://github.com/conda-forge/miniforge/releases
     # Below 2 are included in miniconda base, but not mamba so need to install
     conda-content-trust \
     charset-normalizer \
-
  && /opt/conda/bin/conda clean -ya
 
 RUN conda install -c conda-forge \
@@ -120,10 +118,12 @@ RUN pip install -U --extra-index-url https://pip.repos.neuron.amazonaws.com \
     torch-neuronx==$NEURONX_FRAMEWORK_VERSION \
     transformers-neuronx==$NEURONX_TRANSFORMERS_VERSION \
  && pip install -U "protobuf>=3.18.3,<4" \
+    "transformers==4.45.*" \
     torchserve==${TORCHSERVE_VERSION} \
     torch-model-archiver==${TORCHSERVE_VERSION} \
- && pip install --no-deps --no-cache-dir -U torchvision==0.16.* \
- && pip install --no-deps -U --extra-index-url https://pip.repos.neuron.amazonaws.com neuronx_distributed==$NEURONX_DISTRIBUTED_VERSION
+ && pip install --no-deps --no-cache-dir -U torchvision==0.20.* \
+ && pip install --no-deps -U --extra-index-url https://pip.repos.neuron.amazonaws.com neuronx_distributed==$NEURONX_DISTRIBUTED_VERSION \
+ && pip install -U --extra-index-url https://pip.repos.neuron.amazonaws.com neuronx_distributed_inference==$NEURONX_DISTRIBUTED_INFERENCE_VERSION
 
 RUN useradd -m model-server \
  && mkdir -p /home/model-server/tmp /opt/ml/model \
@@ -166,3 +166,5 @@ EXPOSE 8080 8081
 
 ENTRYPOINT ["python", "/usr/local/bin/dockerd-entrypoint.py"]
 CMD ["/usr/local/bin/entrypoint.sh"]
+
+HEALTHCHECK CMD curl --fail http://localhost:8080/ping || exit 1
\ No newline at end of file
diff --git a/docker/pytorch/inference/2.1.2/Dockerfile.neuronx.cve_allowlist.json b/docker/pytorch/inference/2.5.1/Dockerfile.neuronx.cve_allowlist.json
similarity index 51%
rename from docker/pytorch/inference/2.1.2/Dockerfile.neuronx.cve_allowlist.json
rename to docker/pytorch/inference/2.5.1/Dockerfile.neuronx.cve_allowlist.json
index 536c7dc..30036a2 100644
--- a/docker/pytorch/inference/2.1.2/Dockerfile.neuronx.cve_allowlist.json
+++ b/docker/pytorch/inference/2.5.1/Dockerfile.neuronx.cve_allowlist.json
@@ -1,6 +1,105 @@
 {
-    "CVE-2024-2511": {
-        "description": "Issue summary: Some non-default TLS server configurations can cause unbounded memory growth when processing TLSv1.3 sessions\n\nImpact summary: An attacker may exploit certain server configurations to trigger unbounded memory growth that would lead to a Denial of Service\n\nThis problem can occur in TLSv1.3 if the non-default SSL_OP_NO_TICKET option is being used (but not if early_data support is also configured and the default anti-replay protection is in use). In this case, under certain conditions, the session cache can get into an incorrect state and it will fail to flush properly as it fills. The session cache will continue to grow in an unbounded manner. A malicious client could deliberately create the scenario for this failure to force a Denial of Service. It may also happen by accident in normal operation.\n\nThis issue only affects TLS servers supporting TLSv1.3. It does not affect TLS clients.\n\nThe FIPS modules in 3.2, 3.1 and 3.0 are not affected by this issue. OpenSSL 1.0.2 is also not affected by this ",
+    "CVE-2024-11392": {
+        "description": "Hugging Face Transformers MobileViTV2 Deserialization of Untrusted Data Remote Code Execution Vulnerability. This vulnerability allows remote attackers to execute arbitrary code on affected installations of Hugging Face Transformers. User interaction is required to exploit this vulnerability in that the target must visit a malicious page or open a malicious file.\n\nThe specific flaw exists within the handling of configuration files. The issue results from the lack of proper validation of user-supplied data, which can result in deserialization of untrusted data. An attacker can leverage this vulnerability to execute code in the context of the current user. Was ZDI-CAN-24322.",
+        "remediation": {
+            "recommendation": {
+                "text": "None Provided"
+            }
+        },
+        "score": 8.8,
+        "score_details": {
+            "cvss": {
+                "adjustments": [],
+                "score": 8.8,
+                "scoreSource": "NVD",
+                "scoringVector": "CVSS:3.1/AV:N/AC:L/PR:N/UI:R/S:U/C:H/I:H/A:H",
+                "version": "3.1"
+            }
+        },
+        "severity": "HIGH",
+        "source": "NVD",
+        "source_url": "https://nvd.nist.gov/vuln/detail/CVE-2024-11392",
+        "status": "ACTIVE",
+        "title": "CVE-2024-11392 - transformers",
+        "vulnerability_id": "CVE-2024-11392",
+        "vulnerable_packages": [
+            {
+                "epoch": 0,
+                "filePath": "opt/conda/lib/python3.10/site-packages/transformers-4.46.3.dist-info/METADATA",
+                "name": "transformers",
+                "packageManager": "PYTHONPKG",
+                "version": "4.46.3"
+            }
+        ]
+    },
+    "CVE-2024-11393": {
+        "description": "Hugging Face Transformers MaskFormer Model Deserialization of Untrusted Data Remote Code Execution Vulnerability. This vulnerability allows remote attackers to execute arbitrary code on affected installations of Hugging Face Transformers. User interaction is required to exploit this vulnerability in that the target must visit a malicious page or open a malicious file.\n\nThe specific flaw exists within the parsing of model files. The issue results from the lack of proper validation of user-supplied data, which can result in deserialization of untrusted data. An attacker can leverage this vulnerability to execute code in the context of the current user. Was ZDI-CAN-25191.",
+        "remediation": {
+            "recommendation": {
+                "text": "None Provided"
+            }
+        },
+        "score": 8.8,
+        "score_details": {
+            "cvss": {
+                "adjustments": [],
+                "score": 8.8,
+                "scoreSource": "NVD",
+                "scoringVector": "CVSS:3.1/AV:N/AC:L/PR:N/UI:R/S:U/C:H/I:H/A:H",
+                "version": "3.1"
+            }
+        },
+        "severity": "HIGH",
+        "source": "NVD",
+        "source_url": "https://nvd.nist.gov/vuln/detail/CVE-2024-11393",
+        "status": "ACTIVE",
+        "title": "CVE-2024-11393 - transformers",
+        "vulnerability_id": "CVE-2024-11393",
+        "vulnerable_packages": [
+            {
+                "epoch": 0,
+                "filePath": "opt/conda/lib/python3.10/site-packages/transformers-4.46.3.dist-info/METADATA",
+                "name": "transformers",
+                "packageManager": "PYTHONPKG",
+                "version": "4.46.3"
+            }
+        ]
+    },
+    "CVE-2024-11394": {
+        "description": "Hugging Face Transformers Trax Model Deserialization of Untrusted Data Remote Code Execution Vulnerability. This vulnerability allows remote attackers to execute arbitrary code on affected installations of Hugging Face Transformers. User interaction is required to exploit this vulnerability in that the target must visit a malicious page or open a malicious file.\n\nThe specific flaw exists within the handling of model files. The issue results from the lack of proper validation of user-supplied data, which can result in deserialization of untrusted data. An attacker can leverage this vulnerability to execute code in the context of the current user. Was ZDI-CAN-25012.",
+        "remediation": {
+            "recommendation": {
+                "text": "None Provided"
+            }
+        },
+        "score": 8.8,
+        "score_details": {
+            "cvss": {
+                "adjustments": [],
+                "score": 8.8,
+                "scoreSource": "NVD",
+                "scoringVector": "CVSS:3.1/AV:N/AC:L/PR:N/UI:R/S:U/C:H/I:H/A:H",
+                "version": "3.1"
+            }
+        },
+        "severity": "HIGH",
+        "source": "NVD",
+        "source_url": "https://nvd.nist.gov/vuln/detail/CVE-2024-11394",
+        "status": "ACTIVE",
+        "title": "CVE-2024-11394 - transformers",
+        "vulnerability_id": "CVE-2024-11394",
+        "vulnerable_packages": [
+            {
+                "epoch": 0,
+                "filePath": "opt/conda/lib/python3.10/site-packages/transformers-4.46.3.dist-info/METADATA",
+                "name": "transformers",
+                "packageManager": "PYTHONPKG",
+                "version": "4.46.3"
+            }
+        ]
+    },
+    "CVE-2023-6237": {
+        "description": "Issue summary: Checking excessively long invalid RSA public keys may take a long time.\n\nImpact summary: Applications that use the function EVP_PKEY_public_check() to check RSA public keys may experience long delays. Where the key that is being checked has been obtained from an untrusted source this may lead to a Denial of Service.\n\nWhen function EVP_PKEY_public_check() is called on RSA public keys, a computation is done to confirm that the RSA modulus, n, is composite. For valid RSA keys, n is a product of two or more large primes and this computation completes quickly. However, if n is an overly large prime, then this computation would take a long time.\n\nAn application that calls EVP_PKEY_public_check() and supplies an RSA key obtained from an untrusted source could be vulnerable to a Denial of Service attack.\n\nThe function EVP_PKEY_public_check() is not called from other OpenSSL functions however it is called from the OpenSSL pkey command line application. For that reason that application is also vulnerable",
         "remediation": {
             "recommendation": {
                 "text": "None Provided"
@@ -10,10 +109,10 @@
         "score_details": {},
         "severity": "UNTRIAGED",
         "source": "NVD",
-        "source_url": "https://nvd.nist.gov/vuln/detail/CVE-2024-2511",
+        "source_url": "https://nvd.nist.gov/vuln/detail/CVE-2023-6237",
         "status": "ACTIVE",
-        "title": "CVE-2024-2511 - pyOpenSSL",
-        "vulnerability_id": "CVE-2024-2511",
+        "title": "CVE-2023-6237 - pyOpenSSL",
+        "vulnerability_id": "CVE-2023-6237",
         "vulnerable_packages": [
             {
                 "epoch": 0,
diff --git a/docker/pytorch/training/1.13.1/Dockerfile.neuronx.cve_allowlist.json b/docker/pytorch/training/1.13.1/Dockerfile.neuronx.cve_allowlist.json
deleted file mode 100644
index 44bfaa3..0000000
--- a/docker/pytorch/training/1.13.1/Dockerfile.neuronx.cve_allowlist.json
+++ /dev/null
@@ -1,206 +0,0 @@
-{
-    "CVE-2024-31580": {
-        "description": "PyTorch before v2.2.0 was discovered to contain a heap buffer overflow vulnerability in the component /runtime/vararg_functions.cpp. This vulnerability allows attackers to cause a Denial of Service (DoS) via a crafted input.",
-        "remediation": {
-            "recommendation": {
-                "text": "None Provided"
-            }
-        },
-        "score": 0.0,
-        "score_details": {},
-        "severity": "UNTRIAGED",
-        "source": "NVD",
-        "source_url": "https://nvd.nist.gov/vuln/detail/CVE-2024-31580",
-        "status": "ACTIVE",
-        "title": "CVE-2024-31580 - torch",
-        "vulnerability_id": "CVE-2024-31580",
-        "vulnerable_packages": [
-            {
-                "epoch": 0,
-                "filePath": "usr/local/lib/python3.10/site-packages/torch-1.13.1.dist-info/METADATA",
-                "name": "torch",
-                "packageManager": "PYTHONPKG",
-                "version": "1.13.1"
-            }
-        ]
-    },
-    "CVE-2024-31583": {
-        "description": "Pytorch before version v2.2.0 was discovered to contain a use-after-free vulnerability in torch/csrc/jit/mobile/interpreter.cpp.",
-        "remediation": {
-            "recommendation": {
-                "text": "None Provided"
-            }
-        },
-        "score": 0.0,
-        "score_details": {},
-        "severity": "UNTRIAGED",
-        "source": "NVD",
-        "source_url": "https://nvd.nist.gov/vuln/detail/CVE-2024-31583",
-        "status": "ACTIVE",
-        "title": "CVE-2024-31583 - torch",
-        "vulnerability_id": "CVE-2024-31583",
-        "vulnerable_packages": [
-            {
-                "epoch": 0,
-                "filePath": "usr/local/lib/python3.10/site-packages/torch-1.13.1.dist-info/METADATA",
-                "name": "torch",
-                "packageManager": "PYTHONPKG",
-                "version": "1.13.1"
-            }
-        ]
-    },
-    "CVE-2024-34072": {
-        "description": "sagemaker-python-sdk is a library for training and deploying machine learning models on Amazon SageMaker. The sagemaker.base_deserializers.NumpyDeserializer module before v2.218.0 allows potentially unsafe deserialization when untrusted data is passed as pickled object arrays. This consequently may allow an unprivileged third party to cause remote code execution, denial of service, affecting both confidentiality and integrity. Users are advised to upgrade to version 2.218.0. Users unable to upgrade should not pass pickled numpy object arrays which originated from an untrusted source, or that could have been tampered with. Only pass pickled numpy object arrays from trusted sources.",
-        "remediation": {
-            "recommendation": {
-                "text": "None Provided"
-            }
-        },
-        "score": 0.0,
-        "score_details": {},
-        "severity": "UNTRIAGED",
-        "source": "NVD",
-        "source_url": "https://nvd.nist.gov/vuln/detail/CVE-2024-34072",
-        "status": "ACTIVE",
-        "title": "CVE-2024-34072 - sagemaker",
-        "vulnerability_id": "CVE-2024-34072",
-        "vulnerable_packages": [
-            {
-                "epoch": 0,
-                "filePath": "usr/local/lib/python3.10/site-packages/sagemaker-2.183.0.dist-info/METADATA",
-                "name": "sagemaker",
-                "packageManager": "PYTHONPKG",
-                "version": "2.183.0"
-            }
-        ]
-    },
-    "CVE-2024-34073": {
-        "description": "sagemaker-python-sdk is a library for training and deploying machine learning models on Amazon SageMaker. In affected versions the capture_dependencies function in `sagemaker.serve.save_retrive.version_1_0_0.save.utils` module allows for potentially unsafe Operating System (OS) Command Injection if inappropriate command is passed as the \u201crequirements_path\u201d parameter. This consequently may allow an unprivileged third party to cause remote code execution, denial of service, affecting both confidentiality and integrity. This issue has been addressed in version 2.214.3. Users are advised to upgrade. Users unable to upgrade should not override the \u201crequirements_path\u201d parameter of capture_dependencies function in `sagemaker.serve.save_retrive.version_1_0_0.save.utils`, and instead use the default value.",
-        "remediation": {
-            "recommendation": {
-                "text": "None Provided"
-            }
-        },
-        "score": 0.0,
-        "score_details": {},
-        "severity": "UNTRIAGED",
-        "source": "NVD",
-        "source_url": "https://nvd.nist.gov/vuln/detail/CVE-2024-34073",
-        "status": "ACTIVE",
-        "title": "CVE-2024-34073 - sagemaker",
-        "vulnerability_id": "CVE-2024-34073",
-        "vulnerable_packages": [
-            {
-                "epoch": 0,
-                "filePath": "usr/local/lib/python3.10/site-packages/sagemaker-2.183.0.dist-info/METADATA",
-                "name": "sagemaker",
-                "packageManager": "PYTHONPKG",
-                "version": "2.183.0"
-            }
-        ]
-    },
-    "CVE-2024-3568": {
-        "description": "The huggingface/transformers library is vulnerable to arbitrary code execution through deserialization of untrusted data within the `load_repo_checkpoint()` function of the `TFPreTrainedModel()` class. Attackers can execute arbitrary code and commands by crafting a malicious serialized payload, exploiting the use of `pickle.load()` on data from potentially untrusted sources. This vulnerability allows for remote code execution (RCE) by deceiving victims into loading a seemingly harmless checkpoint during a normal training process, thereby enabling attackers to execute arbitrary code on the targeted machine.",
-        "remediation": {
-            "recommendation": {
-                "text": "None Provided"
-            }
-        },
-        "score": 0.0,
-        "score_details": {},
-        "severity": "UNTRIAGED",
-        "source": "NVD",
-        "source_url": "https://nvd.nist.gov/vuln/detail/CVE-2024-3568",
-        "status": "ACTIVE",
-        "title": "CVE-2024-3568 - transformers, transformers",
-        "vulnerability_id": "CVE-2024-3568",
-        "vulnerable_packages": [
-            {
-                "epoch": 0,
-                "filePath": "usr/local/lib/python3.10/site-packages/transformers-4.36.0.dist-info/METADATA",
-                "name": "transformers",
-                "packageManager": "PYTHONPKG",
-                "version": "4.36.0"
-            },
-            {
-                "epoch": 0,
-                "filePath": "requirements.txt",
-                "name": "transformers",
-                "packageManager": "PIP",
-                "version": "4.36.0"
-            }
-        ]
-    },
-    "CVE-2024-5452": {
-        "description": "A remote code execution (RCE) vulnerability exists in the lightning-ai/pytorch-lightning library version 2.2.1 due to improper handling of deserialized user input and mismanagement of dunder attributes by the `deepdiff` library. The library uses `deepdiff.Delta` objects to modify application state based on frontend actions. However, it is possible to bypass the intended restrictions on modifying dunder attributes, allowing an attacker to construct a serialized delta that passes the deserializer whitelist and contains dunder attributes. When processed, this can be exploited to access other modules, classes, and instances, leading to arbitrary attribute write and total RCE on any self-hosted pytorch-lightning application in its default configuration, as the delta endpoint is enabled by default.",
-        "remediation": {
-            "recommendation": {
-                "text": "None Provided"
-            }
-        },
-        "score": 9.8,
-        "score_details": {
-            "cvss": {
-                "adjustments": [],
-                "score": 9.8,
-                "scoreSource": "NVD",
-                "scoringVector": "CVSS:3.1/AV:N/AC:L/PR:N/UI:N/S:U/C:H/I:H/A:H",
-                "version": "3.1"
-            }
-        },
-        "severity": "CRITICAL",
-        "source": "NVD",
-        "source_url": "https://nvd.nist.gov/vuln/detail/CVE-2024-5452",
-        "status": "ACTIVE",
-        "title": "CVE-2024-5452 - pytorch-lightning, pytorch-lightning",
-        "vulnerability_id": "CVE-2024-5452",
-        "vulnerable_packages": [
-            {
-                "epoch": 0,
-                "filePath": "usr/local/lib/python3.10/site-packages/pytorch_lightning-1.8.6.dist-info/METADATA",
-                "name": "pytorch-lightning",
-                "packageManager": "PYTHONPKG",
-                "version": "1.8.6"
-            },
-            {
-                "epoch": 0,
-                "filePath": "requirements.txt",
-                "name": "pytorch-lightning",
-                "packageManager": "PIP",
-                "version": "1.8.6"
-            }
-        ]
-    },
-    "CVE-2024-5980": {
-        "description": "A vulnerability in the /v1/runs API endpoint of lightning-ai/pytorch-lightning v2.2.4 allows attackers to exploit path traversal when extracting tar.gz files. When the LightningApp is running with the plugin_server, attackers can deploy malicious tar.gz plugins that embed arbitrary files with path traversal vulnerabilities. This can result in arbitrary files being written to any directory in the victim's local file system, potentially leading to remote code execution.",
-        "remediation": {
-            "recommendation": {
-                "text": "None Provided"
-            }
-        },
-        "score": 0.0,
-        "score_details": {},
-        "severity": "UNTRIAGED",
-        "source": "NVD",
-        "source_url": "https://nvd.nist.gov/vuln/detail/CVE-2024-5980",
-        "status": "ACTIVE",
-        "title": "CVE-2024-5980 - pytorch-lightning, pytorch-lightning",
-        "vulnerability_id": "CVE-2024-5980",
-        "vulnerable_packages": [
-            {
-                "epoch": 0,
-                "filePath": "usr/local/lib/python3.10/site-packages/pytorch_lightning-1.8.6.dist-info/METADATA",
-                "name": "pytorch-lightning",
-                "packageManager": "PYTHONPKG",
-                "version": "1.8.6"
-            },
-            {
-                "epoch": 0,
-                "filePath": "requirements.txt",
-                "name": "pytorch-lightning",
-                "packageManager": "PIP",
-                "version": "1.8.6"
-            }
-        ]
-    }
-}
\ No newline at end of file
diff --git a/docker/pytorch/training/2.1.2/Dockerfile.neuronx b/docker/pytorch/training/2.5.1/Dockerfile.neuronx
similarity index 93%
rename from docker/pytorch/training/2.1.2/Dockerfile.neuronx
rename to docker/pytorch/training/2.5.1/Dockerfile.neuronx
index 9a3a0c7..414c2f1 100644
--- a/docker/pytorch/training/2.1.2/Dockerfile.neuronx
+++ b/docker/pytorch/training/2.5.1/Dockerfile.neuronx
@@ -1,16 +1,16 @@
-FROM public.ecr.aws/docker/library/ubuntu:20.04
+FROM public.ecr.aws/docker/library/ubuntu:22.04
 
 LABEL maintainer="Amazon AI"
 LABEL dlc_major_version="1"
 
 # Neuron SDK components version numbers
-ARG NEURONX_DISTRIBUTED_VERSION=0.9.0
-ARG NEURONX_DISTRIBUTED_TRAINING_VERSION=1.0.1
-ARG NEURONX_CC_VERSION=2.15.143.0
-ARG NEURONX_FRAMEWORK_VERSION=2.1.2.2.3.2
-ARG NEURONX_COLLECTIVES_LIB_VERSION=2.22.33.0-d2128d1aa
-ARG NEURONX_RUNTIME_LIB_VERSION=2.22.19.0-5856c0b42
-ARG NEURONX_TOOLS_VERSION=2.19.0.0
+ARG NEURONX_DISTRIBUTED_VERSION=0.10.0
+ARG NEURONX_DISTRIBUTED_TRAINING_VERSION=1.1.0
+ARG NEURONX_CC_VERSION=2.16.345.0
+ARG NEURONX_FRAMEWORK_VERSION=2.5.1.2.4.0
+ARG NEURONX_COLLECTIVES_LIB_VERSION=2.23.133.0-3e70920f2
+ARG NEURONX_RUNTIME_LIB_VERSION=2.23.110.0-9b5179492
+ARG NEURONX_TOOLS_VERSION=2.20.204.0
 
 ARG PYTHON=python3.10
 ARG PYTHON_VERSION=3.10.12
@@ -32,7 +32,7 @@ ENV LD_LIBRARY_PATH="${LD_LIBRARY_PATH}:/opt/amazon/efa/lib"
 ENV LD_LIBRARY_PATH="${LD_LIBRARY_PATH}:/opt/amazon/efa/lib64"
 ENV LD_LIBRARY_PATH="${LD_LIBRARY_PATH}:/opt/amazon/openmpi/lib64"
 ENV LD_LIBRARY_PATH="${LD_LIBRARY_PATH}:/usr/local/lib"
-ENV PATH /opt/aws/neuron/bin/:$PATH
+ENV PATH="/opt/aws/neuron/bin:${PATH}"
 ENV SAGEMAKER_TRAINING_MODULE=sagemaker_pytorch_container.training:main
 ENV DGLBACKEND=pytorch
 
@@ -45,34 +45,35 @@ RUN apt-get update \
     curl \
     emacs \
     git \
+    gnupg2 \
+    gpg-agent \
     jq \
     libopencv-dev \
-    openjdk-8-jdk-headless \
-    openjdk-8-jdk \
-    openjdk-8-jre \
     libglib2.0-0 \
     libgl1-mesa-glx \
     libsm6 \
     libxext6 \
     libxrender-dev \
-    openjdk-11-jdk \
-    software-properties-common \
-    wget \
-    unzip \
-    vim \
-    zlib1g-dev \
-    openssl \
     libssl-dev \
     libsqlite3-dev \
     libgdbm-dev \
     libc6-dev \
     libbz2-dev \
     libncurses-dev \
-    tk-dev \
     libffi-dev \
     libcap-dev \
-    gnupg2 \
-    gpg-agent \
+    libhwloc-dev \
+    openjdk-8-jdk-headless \
+    openjdk-8-jdk \
+    openjdk-8-jre \
+    openjdk-11-jdk \
+    openssl \
+    software-properties-common \
+    tk-dev \
+    unzip \
+    wget \
+    vim \
+    zlib1g-dev \
  && rm -rf /var/lib/apt/lists/* \
  && apt-get clean
 
@@ -165,7 +166,7 @@ RUN wget https://raw.githubusercontent.com/aws-neuron/neuronx-distributed-traini
     && ${PIP} install --no-cache-dir -r requirements.txt /root/apex/dist/apex-0.1-py3-none-any.whl \
     && ${PIP} install --force-reinstall "multiprocess==0.70.16" \
     "dill==0.3.8" \
-    "torch==2.1.2"
+    "torch==2.5.1"
 
 
 RUN ${PIP} install --force-reinstall --no-deps neuronx_distributed_training==$NEURONX_DISTRIBUTED_TRAINING_VERSION --extra-index-url https://pip.repos.neuron.amazonaws.com
@@ -215,7 +216,7 @@ RUN rm -rf /var/lib/apt/lists/* \
 # torchvision needed for MLP. since it depends on torch and torch neuron/torch
 # is already installed install it with nodeps
 RUN pip3 install --no-cache-dir --no-deps -U \
-    torchvision==0.16.*
+    torchvision==0.20.*
 
 # Needed for running bert training scripts
 RUN pip3 install --no-cache-dir -U \
@@ -241,8 +242,10 @@ RUN HOME_DIR=/root \
  && rm -rf ${HOME_DIR}/oss_compliance* \
  && rm -rf /tmp/tmp*
 
-RUN curl -o /license.txt  https://aws-dlc-licenses.s3.amazonaws.com/pytorch-2.1/license.txt
+RUN curl -o /license.txt  https://aws-dlc-licenses.s3.amazonaws.com/pytorch-2.5/license.txt
 
 # Starts framework
 ENTRYPOINT ["bash", "-m", "start_with_right_hostname.sh"]
 CMD ["/bin/bash"]
+
+HEALTHCHECK CMD curl --fail http://localhost:8080/ping || exit 1
diff --git a/docker/pytorch/training/2.1.2/Dockerfile.neuronx.cve_allowlist.json b/docker/pytorch/training/2.5.1/Dockerfile.neuronx.cve_allowlist.json
similarity index 68%
rename from docker/pytorch/training/2.1.2/Dockerfile.neuronx.cve_allowlist.json
rename to docker/pytorch/training/2.5.1/Dockerfile.neuronx.cve_allowlist.json
index ae57548..31e2487 100644
--- a/docker/pytorch/training/2.1.2/Dockerfile.neuronx.cve_allowlist.json
+++ b/docker/pytorch/training/2.5.1/Dockerfile.neuronx.cve_allowlist.json
@@ -32,6 +32,126 @@
             }
         ]
     },
+    "CVE-2024-11392": {
+        "description": "Hugging Face Transformers MobileViTV2 Deserialization of Untrusted Data Remote Code Execution Vulnerability. This vulnerability allows remote attackers to execute arbitrary code on affected installations of Hugging Face Transformers. User interaction is required to exploit this vulnerability in that the target must visit a malicious page or open a malicious file.\n\nThe specific flaw exists within the handling of configuration files. The issue results from the lack of proper validation of user-supplied data, which can result in deserialization of untrusted data. An attacker can leverage this vulnerability to execute code in the context of the current user. Was ZDI-CAN-24322.",
+        "remediation": {
+            "recommendation": {
+                "text": "None Provided"
+            }
+        },
+        "score": 8.8,
+        "score_details": {
+            "cvss": {
+                "adjustments": [],
+                "score": 8.8,
+                "scoreSource": "NVD",
+                "scoringVector": "CVSS:3.1/AV:N/AC:L/PR:N/UI:R/S:U/C:H/I:H/A:H",
+                "version": "3.1"
+            }
+        },
+        "severity": "HIGH",
+        "source": "NVD",
+        "source_url": "https://nvd.nist.gov/vuln/detail/CVE-2024-11392",
+        "status": "ACTIVE",
+        "title": "CVE-2024-11392 - transformers, transformers",
+        "vulnerability_id": "CVE-2024-11392",
+        "vulnerable_packages": [
+            {
+                "epoch": 0,
+                "filePath": "usr/local/lib/python3.10/site-packages/transformers-4.36.0.dist-info/METADATA",
+                "name": "transformers",
+                "packageManager": "PYTHONPKG",
+                "version": "4.36.0"
+            },
+            {
+                "epoch": 0,
+                "filePath": "requirements.txt",
+                "name": "transformers",
+                "packageManager": "PIP",
+                "version": "4.36.0"
+            }
+        ]
+    },
+    "CVE-2024-11393": {
+        "description": "Hugging Face Transformers MaskFormer Model Deserialization of Untrusted Data Remote Code Execution Vulnerability. This vulnerability allows remote attackers to execute arbitrary code on affected installations of Hugging Face Transformers. User interaction is required to exploit this vulnerability in that the target must visit a malicious page or open a malicious file.\n\nThe specific flaw exists within the parsing of model files. The issue results from the lack of proper validation of user-supplied data, which can result in deserialization of untrusted data. An attacker can leverage this vulnerability to execute code in the context of the current user. Was ZDI-CAN-25191.",
+        "remediation": {
+            "recommendation": {
+                "text": "None Provided"
+            }
+        },
+        "score": 8.8,
+        "score_details": {
+            "cvss": {
+                "adjustments": [],
+                "score": 8.8,
+                "scoreSource": "NVD",
+                "scoringVector": "CVSS:3.1/AV:N/AC:L/PR:N/UI:R/S:U/C:H/I:H/A:H",
+                "version": "3.1"
+            }
+        },
+        "severity": "HIGH",
+        "source": "NVD",
+        "source_url": "https://nvd.nist.gov/vuln/detail/CVE-2024-11393",
+        "status": "ACTIVE",
+        "title": "CVE-2024-11393 - transformers, transformers",
+        "vulnerability_id": "CVE-2024-11393",
+        "vulnerable_packages": [
+            {
+                "epoch": 0,
+                "filePath": "usr/local/lib/python3.10/site-packages/transformers-4.36.0.dist-info/METADATA",
+                "name": "transformers",
+                "packageManager": "PYTHONPKG",
+                "version": "4.36.0"
+            },
+            {
+                "epoch": 0,
+                "filePath": "requirements.txt",
+                "name": "transformers",
+                "packageManager": "PIP",
+                "version": "4.36.0"
+            }
+        ]
+    },
+    "CVE-2024-11394": {
+        "description": "Hugging Face Transformers Trax Model Deserialization of Untrusted Data Remote Code Execution Vulnerability. This vulnerability allows remote attackers to execute arbitrary code on affected installations of Hugging Face Transformers. User interaction is required to exploit this vulnerability in that the target must visit a malicious page or open a malicious file.\n\nThe specific flaw exists within the handling of model files. The issue results from the lack of proper validation of user-supplied data, which can result in deserialization of untrusted data. An attacker can leverage this vulnerability to execute code in the context of the current user. Was ZDI-CAN-25012.",
+        "remediation": {
+            "recommendation": {
+                "text": "None Provided"
+            }
+        },
+        "score": 8.8,
+        "score_details": {
+            "cvss": {
+                "adjustments": [],
+                "score": 8.8,
+                "scoreSource": "NVD",
+                "scoringVector": "CVSS:3.1/AV:N/AC:L/PR:N/UI:R/S:U/C:H/I:H/A:H",
+                "version": "3.1"
+            }
+        },
+        "severity": "HIGH",
+        "source": "NVD",
+        "source_url": "https://nvd.nist.gov/vuln/detail/CVE-2024-11394",
+        "status": "ACTIVE",
+        "title": "CVE-2024-11394 - transformers, transformers",
+        "vulnerability_id": "CVE-2024-11394",
+        "vulnerable_packages": [
+            {
+                "epoch": 0,
+                "filePath": "usr/local/lib/python3.10/site-packages/transformers-4.36.0.dist-info/METADATA",
+                "name": "transformers",
+                "packageManager": "PYTHONPKG",
+                "version": "4.36.0"
+            },
+            {
+                "epoch": 0,
+                "filePath": "requirements.txt",
+                "name": "transformers",
+                "packageManager": "PIP",
+                "version": "4.36.0"
+            }
+        ]
+    },
     "CVE-2024-31580": {
         "description": "PyTorch before v2.2.0 was discovered to contain a heap buffer overflow vulnerability in the component /runtime/vararg_functions.cpp. This vulnerability allows attackers to cause a Denial of Service (DoS) via a crafted input.",
         "remediation": {