diff --git a/.github/workflows/cuda/Linux.sh b/.github/workflows/cuda/Linux.sh
index 798b469..74086fc 100644
--- a/.github/workflows/cuda/Linux.sh
+++ b/.github/workflows/cuda/Linux.sh
@@ -1,16 +1,16 @@
 #!/bin/bash
 
+# Ubuntu version
+UBUNTU_VER=${1}
+OS=$(echo ${UBUNTU_VER} | tr -d '-' | tr -d '.')
+
 # CUDA major and minor version
-CUDA_VER_FULL=${1}
+CUDA_VER_FULL=${2}
 CUDA_VER_ARR=($(echo ${CUDA_VER_FULL} | tr "." " "))
 CUDA_VER="${CUDA_VER_ARR[0]}.${CUDA_VER_ARR[1]}"
 CUDA_VER_ID="${CUDA_VER_ARR[0]}_${CUDA_VER_ARR[1]}"
 CUDA_VER_SHORT="cu${CUDA_VER_ARR[0]}${CUDA_VER_ARR[1]}"
 
-# Took from https://github.com/pyg-team/pyg-lib/
-
-OS=ubuntu2004
-
 case ${CUDA_VER_SHORT} in
   cu121)
     CUDA=12.1
diff --git a/.github/workflows/docker.yml b/.github/workflows/docker.yml
index 77d047f..8a99e68 100644
--- a/.github/workflows/docker.yml
+++ b/.github/workflows/docker.yml
@@ -11,15 +11,16 @@ on:
 
 jobs:
   build:
-    name: ${{ matrix.os }}-cuda-${{ matrix.cuda-version }}-torch-${{ matrix.torch-version }}-${{ matrix.cmake-build-type }}
+    name: ${{ matrix.ubuntu-version }}-cuda-${{ matrix.cuda-version }}-torch-${{ matrix.torch-version }}-${{ matrix.cmake-build-type }}
     runs-on: ${{ matrix.os }}
     strategy:
       fail-fast: false
       matrix:
         os: [ubuntu-22.04] # [ubuntu-22.04, ubuntu-20.04, ubuntu-18.04]
         arch: [x64] # [x64, x86]
-        torch-version: [2.1.2, 2.2.1] # [1.12.0, 1.13.0, 2.0.0, 2.1.0, 2.1.1, 2.1.2, 2.2.0, 2.2.1]
-        cuda-version: [11.8.0, 12.1.1] # [12.3.1, 12.1.1, 11.8.0, 11.7.1, 11.6.2, 11.5.2,11.4.4, 11.3.1, 11.2.2, 11.1.1, 11.0.3, cpu]
+        ubuntu-version: [22.04, 20.04] # [18.04, 20.04, 22.04]
+        torch-version: [2.1.2] # [1.12.0, 1.13.0, 2.0.0, 2.1.0, 2.1.1, 2.1.2, 2.2.0, 2.2.1]
+        cuda-version: [11.8.0] # [12.3.1, 12.1.1, 11.8.0, 11.7.1, 11.6.2, 11.5.2,11.4.4, 11.3.1, 11.2.2, 11.1.1, 11.0.3, cpu]
         cmake-build-type: [Release] # [Debug, ClangTidy]
     steps:
       - name: Checkout Repository
@@ -33,6 +34,7 @@ jobs:
           context: .
           file: ./Dockerfile
           build-args: |
+            UBUNTU_VERSION=${{ matrix.ubuntu-version }}
             CUDA_VERSION=${{ matrix.cuda-version }}
             TORCH_VERSION=${{ matrix.torch-version }}
             CMAKE_BUILD_TYPE=${{ matrix.cmake-build-type }}
diff --git a/.github/workflows/hip.yml b/.github/workflows/hip.yml
index 962631d..6b2de65 100644
--- a/.github/workflows/hip.yml
+++ b/.github/workflows/hip.yml
@@ -93,7 +93,7 @@ jobs:
       - name: Install CUDA
         if: ${{ steps.cuda-cache.outputs.cache-hit != 'true' }}
         run: |
-          bash .github/workflows/cuda/${{ runner.os }}.sh ${{ matrix.cuda-version }}
+          bash .github/workflows/cuda/${{ runner.os }}.sh ${{ matrix.os }} ${{ matrix.cuda-version }}
 
       - name: Restore ROCm Cache
         uses: actions/cache@v4
diff --git a/.github/workflows/ubuntu.yml b/.github/workflows/ubuntu.yml
index fcb91b4..973278f 100644
--- a/.github/workflows/ubuntu.yml
+++ b/.github/workflows/ubuntu.yml
@@ -66,7 +66,7 @@ jobs:
       - name: Install CUDA
         if: ${{ steps.cuda-cache.outputs.cache-hit != 'true' }}
         run: |
-          bash .github/workflows/cuda/${{ runner.os }}.sh ${{ matrix.cuda-version }}
+          bash .github/workflows/cuda/${{ runner.os }}.sh ${{ matrix.os }} ${{ matrix.cuda-version }}
 
       - name: Setup CUDA
         shell: bash
diff --git a/Dockerfile b/Dockerfile
index 8713073..fdd7182 100644
--- a/Dockerfile
+++ b/Dockerfile
@@ -2,6 +2,7 @@ ARG UBUNTU_VERSION=22.04
 
 FROM ubuntu:${UBUNTU_VERSION}
 
+ARG UBUNTU_VERSION
 ARG TORCH_VERSION=2.2.1
 ARG CUDA_VERSION=12.1.1
 ARG TORCH_CUDA_ARCH_LIST=7.0;7.5
@@ -18,6 +19,17 @@ WORKDIR /code
 # Copy everything
 COPY . ./
 
+# Upgrade cmake if Ubuntu version is 20.04
+RUN if [[ "$UBUNTU_VERSION" = "20.04" ]]; then \
+        apt-get update && \
+        apt-get install -y ca-certificates gpg wget && \
+        wget -O - https://apt.kitware.com/keys/kitware-archive-latest.asc 2>/dev/null | gpg --dearmor - | tee /usr/share/keyrings/kitware-archive-keyring.gpg >/dev/null && \
+        echo 'deb [signed-by=/usr/share/keyrings/kitware-archive-keyring.gpg] https://apt.kitware.com/ubuntu/ focal main' | tee /etc/apt/sources.list.d/kitware.list >/dev/null && \
+        apt-get update && \
+        apt-get install kitware-archive-keyring && \
+        echo 'deb [signed-by=/usr/share/keyrings/kitware-archive-keyring.gpg] https://apt.kitware.com/ubuntu/ focal-rc main' | tee -a /etc/apt/sources.list.d/kitware.list >/dev/null; \
+    fi
+
 # Install build dependencies
 RUN apt-get update && \
     apt-get install -y \
@@ -35,7 +47,7 @@ RUN apt-get update && \
 
 
 # Install CUDA
-RUN bash .github/workflows/cuda/Linux.sh ${CUDA_VERSION}
+RUN bash .github/workflows/cuda/Linux.sh "ubuntu-${UBUNTU_VERSION}" ${CUDA_VERSION}
 
 # Install libtorch
 RUN wget --no-check-certificate -nv https://download.pytorch.org/libtorch/cu"${CUDA_VERSION%%.*}"$(echo $CUDA_VERSION | cut -d'.' -f2)/libtorch-cxx11-abi-shared-with-deps-${TORCH_VERSION}%2Bcu"${CUDA_VERSION%%.*}"$(echo $CUDA_VERSION | cut -d'.' -f2).zip -O libtorch.zip && \