Upgrade cmake version (ubuntu-20.04)

Dockerfile

@@ -2,7 +2,7 @@ ARG UBUNTU_VERSION=22.04
 
 FROM ubuntu:${UBUNTU_VERSION}
 
-ARG UBUNTU_VERSION=22.04
+ARG UBUNTU_VERSION
 ARG TORCH_VERSION=2.2.1
 ARG CUDA_VERSION=12.1.1
 ARG TORCH_CUDA_ARCH_LIST=7.0;7.5
@@ -19,6 +19,11 @@ WORKDIR /code
 # Copy everything
 COPY . ./
 
+# Upgrade cmake if Ubuntu version is 20.04
+RUN if [[ "$UBUNTU_VERSION" = "20.04" ]]; then \
+        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 \