diff --git a/.github/workflows/docker.yml b/.github/workflows/docker.yml
index efc85a8..64ea3ce 100644
--- a/.github/workflows/docker.yml
+++ b/.github/workflows/docker.yml
@@ -21,9 +21,9 @@ jobs:
     strategy:
       fail-fast: false
       matrix:
+        image: ["ubuntu:24.04", "nvidia/cuda:12.6.2-base-ubuntu24.04"]
         # Must match version at https://www.python.org/ftp/python/
-        python: ["3.10.13"]
-        system: ["cpu-headless", "gpu-headless"]
+        python: [3.10.13]
 
     # Sets the permissions granted to the `GITHUB_TOKEN` for the actions in this job.
     permissions:
@@ -41,6 +41,12 @@ jobs:
           username: ${{ github.actor }}
           password: ${{ secrets.GITHUB_TOKEN }}
 
+      - name: Sanitize docker tag
+        run:
+          PREFIX_DOCKER_TAG="OnnxTR-${{ matrix.image == 'ubuntu:24.04' && 'cpu' || 'gpu' }}-py${{ matrix.python }}-"
+          PREFIX_DOCKER_TAG=$(echo PREFIX_DOCKER_TAG|sed 's/,/-/g')
+          echo PREFIX_DOCKER_TAG=${PREFIX_DOCKER_TAG} >> $GITHUB_ENV
+
       - name: Extract metadata (tags, labels) for Docker
         id: meta
         uses: docker/metadata-action@v5
@@ -48,18 +54,21 @@ jobs:
           images: ${{ env.REGISTRY }}/${{ github.repository }}
           tags: |
             # used only on schedule event
-            type=schedule,pattern={{date 'YYYY-MM'}},prefix=OnnxTR-${{ matrix.system }}-py${{ matrix.python }}-
+            type=schedule,pattern={{date 'YYYY-MM'}},prefix=${{ env.PREFIX_DOCKER_TAG }}
             # used only if a tag following semver is published
-            type=semver,pattern={{raw}},prefix=OnnxTR-${{ matrix.system }}-py${{ matrix.python }}-
+            type=semver,pattern={{raw}},prefix=${{ env.PREFIX_DOCKER_TAG }}
 
       - name: Build Docker image
         id: build
         uses: docker/build-push-action@v6
         with:
           context: .
+          # Set SYSTEM to 'cpu-headless' if the image is ubuntu:24.04
+          # and 'gpu-headless' otherwise
           build-args: |
+            BASE_IMAGE=${{ matrix.image }}
+            SYSTEM=${{ matrix.image == 'ubuntu:24.04' && 'cpu-headless' || 'gpu-headless' }}
             PYTHON_VERSION=${{ matrix.python }}
-            SYSTEM=${{ matrix.system }}
             ONNXTR_REPO=${{ github.repository }}
             ONNXTR_VERSION=${{ github.sha }}
           push: false  # push only if `import onnxtr` works
@@ -70,13 +79,16 @@ jobs:
 
       - name: Push Docker image
         # Push only if the CI is not triggered by "PR on main"
-        if: github.ref == 'refs/heads/main' && github.event_name != 'pull_request'
+        if: ${{ (github.ref == 'refs/heads/main' && github.event_name != 'pull_request') || (startsWith(github.ref, 'refs/tags') && github.event_name == 'push') }}
         uses: docker/build-push-action@v6
         with:
           context: .
+          # Set SYSTEM to 'cpu-headless' if the image is ubuntu:24.04
+          # and 'gpu-headless' otherwise
           build-args: |
+            BASE_IMAGE=${{ matrix.image }}
+            SYSTEM=${{ matrix.image == 'ubuntu:24.04' && 'cpu-headless' || 'gpu-headless' }}
             PYTHON_VERSION=${{ matrix.python }}
-            SYSTEM=${{ matrix.system }}
             ONNXTR_REPO=${{ github.repository }}
             ONNXTR_VERSION=${{ github.sha }}
           push: true
diff --git a/Dockerfile b/Dockerfile
index a4ff3c5..4813303 100644
--- a/Dockerfile
+++ b/Dockerfile
@@ -1,40 +1,14 @@
-FROM ubuntu:22.04
+ARG BASE_IMAGE
+
+FROM ${BASE_IMAGE}
 
 ENV DEBIAN_FRONTEND=noninteractive
 ENV LANG=C.UTF-8
 ENV PYTHONUNBUFFERED=1
 ENV PYTHONDONTWRITEBYTECODE=1
 
-ARG SYSTEM=gpu
-
-# Enroll NVIDIA GPG public key and install CUDA
-RUN if [ "$SYSTEM" = "gpu" ]; then \
-    apt-get update && \
-    apt-get install -y gnupg ca-certificates wget && \
-    # - Install Nvidia repo keys
-    # - See: https://docs.nvidia.com/cuda/cuda-installation-guide-linux/index.html#network-repo-installation-for-ubuntu
-    wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2204/x86_64/cuda-keyring_1.1-1_all.deb && \
-    dpkg -i cuda-keyring_1.1-1_all.deb && \
-    apt-get update && apt-get install -y --no-install-recommends \
-    # NOTE: The following CUDA_VERSION, CUDNN_VERSION, and NVINFER_VERSION are for CUDA 11.8
-    # - this needs to match exactly with the host system otherwise the onnxruntime-gpu package isn't able to work correct. !!
-    cuda-command-line-tools-11-8 \
-    cuda-cudart-dev-11-8 \
-    cuda-nvcc-11-8 \
-    cuda-cupti-11-8 \
-    cuda-nvprune-11-8 \
-    cuda-libraries-11-8 \
-    cuda-nvrtc-11-8 \
-    libcufft-11-8 \
-    libcurand-11-8 \
-    libcusolver-11-8 \
-    libcusparse-11-8 \
-    libcublas-11-8 \
-    # - CuDNN: https://docs.nvidia.com/deeplearning/sdk/cudnn-install/index.html#ubuntu-network-installation
-    libcudnn8=8.6.0.163-1+cuda11.8 \
-    libnvinfer-plugin8=8.6.1.6-1+cuda11.8 \
-    libnvinfer8=8.6.1.6-1+cuda11.8; \
-fi
+ARG SYSTEM
+ARG PYTHON_VERSION
 
 RUN apt-get update && apt-get install -y --no-install-recommends \
     # - Other packages
@@ -54,7 +28,6 @@ RUN apt-get update && apt-get install -y --no-install-recommends \
 fi
 
 # Install Python
-ARG PYTHON_VERSION=3.10.13
 
 RUN wget http://www.python.org/ftp/python/$PYTHON_VERSION/Python-$PYTHON_VERSION.tgz && \
     tar -zxf Python-$PYTHON_VERSION.tgz && \