From a6d0e97f84aa88f8dc0011e3d6c76d14433e289f Mon Sep 17 00:00:00 2001
From: Anton Zamyatin <e12223389@student.tuwien.ac.at>
Date: Sun, 2 Mar 2025 14:09:05 +0100
Subject: [PATCH 01/16] :whale: Add Dockerfile

---
 Dockerfile | 23 +++++++++++++++++++++++
 1 file changed, 23 insertions(+)
 create mode 100644 Dockerfile

diff --git a/Dockerfile b/Dockerfile
new file mode 100644
index 00000000..9d1b118c
--- /dev/null
+++ b/Dockerfile
@@ -0,0 +1,23 @@
+FROM pytorch/pytorch:2.6.0-cuda12.6-cudnn9-runtime
+
+ENV DEBIAN_FRONTEND=noninteractive \
+    PYTHONUNBUFFERED=1 \
+    LC_ALL=C.UTF-8 \
+    LANG=C.UTF-8
+
+RUN apt-get update && apt-get install -y \
+    git \
+    curl \
+    wget \
+    vim \
+    python3-pip \
+    python3-venv \
+    && rm -rf /var/lib/apt/lists/*
+
+WORKDIR /workspace
+
+COPY pyproject.toml .
+RUN pip install --no-cache-dir .
+
+EXPOSE 8888
+CMD [ "/bin/bash" ]

From 8b416c0be18f88d3af521002f21cacb7fd2832c2 Mon Sep 17 00:00:00 2001
From: Anton Zamyatin <e12223389@student.tuwien.ac.at>
Date: Sun, 2 Mar 2025 18:10:08 +0100
Subject: [PATCH 02/16] :whale: Fix flit build errors

---
 Dockerfile | 13 +++++++++----
 1 file changed, 9 insertions(+), 4 deletions(-)

diff --git a/Dockerfile b/Dockerfile
index 9d1b118c..9bb48918 100644
--- a/Dockerfile
+++ b/Dockerfile
@@ -7,17 +7,22 @@ ENV DEBIAN_FRONTEND=noninteractive \
 
 RUN apt-get update && apt-get install -y \
     git \
-    curl \
-    wget \
-    vim \
     python3-pip \
-    python3-venv \
     && rm -rf /var/lib/apt/lists/*
 
 WORKDIR /workspace
 
+# Copy README.md and torch_uncertainy module (required by pyproject.toml, othwise flit build will fail)
+COPY README.md .
+COPY torch_uncertainty ./torch_uncertainty
+
+# Copy dependency file
 COPY pyproject.toml .
+
+# Install dependencies
 RUN pip install --no-cache-dir .
 
+# Expose port 8888 for TensorBoard and Jupyter Notebook
 EXPOSE 8888
+
 CMD [ "/bin/bash" ]

From 3ef913257a71f4cc366cecfaae330a5a1cf6f666 Mon Sep 17 00:00:00 2001
From: Anton Zamyatin <e12223389@student.tuwien.ac.at>
Date: Sun, 2 Mar 2025 18:25:20 +0100
Subject: [PATCH 03/16] :whale: Install OpenSSH and start server when container
 starts

---
 Dockerfile | 15 ++++++++++++++-
 1 file changed, 14 insertions(+), 1 deletion(-)

diff --git a/Dockerfile b/Dockerfile
index 9bb48918..cfb550bb 100644
--- a/Dockerfile
+++ b/Dockerfile
@@ -5,6 +5,7 @@ ENV DEBIAN_FRONTEND=noninteractive \
     LC_ALL=C.UTF-8 \
     LANG=C.UTF-8
 
+# Install git and pip
 RUN apt-get update && apt-get install -y \
     git \
     python3-pip \
@@ -22,7 +23,19 @@ COPY pyproject.toml .
 # Install dependencies
 RUN pip install --no-cache-dir .
 
+# Install OpenSSH Server
+RUN apt-get update && apt-get install -y openssh-server && rm -rf /var/lib/apt/lists/*
+
+# Create SSH directory & keys
+RUN mkdir -p /var/run/sshd && echo 'root:root' | chpasswd
+
+# Allow root login via SSH
+RUN sed -i 's/#PermitRootLogin prohibit-password/PermitRootLogin yes/' /etc/ssh/sshd_config
+
 # Expose port 8888 for TensorBoard and Jupyter Notebook
 EXPOSE 8888
+# Expose port 22 for SSH
+EXPOSE 22
 
-CMD [ "/bin/bash" ]
+# Ensure the SSH server starts on container launch
+CMD ["/usr/sbin/sshd", "-D"]

From 4ffeb34df21d4add9ac54b16646e66a4639ccf80 Mon Sep 17 00:00:00 2001
From: Anton Zamyatin <e12223389@student.tuwien.ac.at>
Date: Mon, 3 Mar 2025 14:05:31 +0100
Subject: [PATCH 04/16] :whale: Fix SSH key not being recognized and add shell
 prompt customization

---
 Dockerfile | 64 ++++++++++++++++++++++++++++++++++--------------------
 1 file changed, 41 insertions(+), 23 deletions(-)

diff --git a/Dockerfile b/Dockerfile
index cfb550bb..4ae9edca 100644
--- a/Dockerfile
+++ b/Dockerfile
@@ -5,37 +5,55 @@ ENV DEBIAN_FRONTEND=noninteractive \
     LC_ALL=C.UTF-8 \
     LANG=C.UTF-8
 
-# Install git and pip
+# Install Git and OpenSSH Server (PyTorch's base image alraedy includes Conda and Pip)
 RUN apt-get update && apt-get install -y \
     git \
-    python3-pip \
+    openssh-server \
     && rm -rf /var/lib/apt/lists/*
 
 WORKDIR /workspace
 
 # Copy README.md and torch_uncertainy module (required by pyproject.toml, othwise flit build will fail)
-COPY README.md .
-COPY torch_uncertainty ./torch_uncertainty
+COPY README.md /workspace/
+COPY torch_uncertainty /workspace/torch_uncertainty
 
 # Copy dependency file
-COPY pyproject.toml .
+COPY pyproject.toml /workspace/
 
 # Install dependencies
-RUN pip install --no-cache-dir .
-
-# Install OpenSSH Server
-RUN apt-get update && apt-get install -y openssh-server && rm -rf /var/lib/apt/lists/*
-
-# Create SSH directory & keys
-RUN mkdir -p /var/run/sshd && echo 'root:root' | chpasswd
-
-# Allow root login via SSH
-RUN sed -i 's/#PermitRootLogin prohibit-password/PermitRootLogin yes/' /etc/ssh/sshd_config
-
-# Expose port 8888 for TensorBoard and Jupyter Notebook
-EXPOSE 8888
-# Expose port 22 for SSH
-EXPOSE 22
-
-# Ensure the SSH server starts on container launch
-CMD ["/usr/sbin/sshd", "-D"]
+RUN pip install --no-cache-dir ".[all]"
+
+# Always activate Conda when opening a new terminal
+RUN echo "source /opt/conda/bin/activate" >> /root/.bashrc
+
+# Customize the Bash prompt
+RUN echo 'force_color_prompt=yes' >> /root/.bashrc && \
+    echo 'PS1="\[\033[01;34m\]\W\[\033[00m\]\$ "' >> /root/.bashrc && \
+    echo 'if [ -x /usr/bin/dircolors ]; then' >> /root/.bashrc && \
+    echo '    test -r ~/.dircolors && eval "$(dircolors -b ~/.dircolors)" || eval "$(dircolors -b)"' >> /root/.bashrc && \
+    echo '    alias ls="ls --color=auto"' >> /root/.bashrc && \
+    echo '    alias grep="grep --color=auto"' >> /root/.bashrc && \
+    echo '    alias fgrep="fgrep --color=auto"' >> /root/.bashrc && \
+    echo '    alias egrep="egrep --color=auto"' >> /root/.bashrc && \
+    echo '    cd /workspace' >> /root/.bashrc && \
+    echo 'fi' >> /root/.bashrc
+
+# Allow SSH login as root without a password (key-based authentication only)
+RUN echo "PermitRootLogin yes" >> /etc/ssh/sshd_config && \
+    echo "PubkeyAuthentication yes" >> /etc/ssh/sshd_config && \
+    echo "AuthorizedKeysFile .ssh/authorized_keys" >> /etc/ssh/sshd_config
+
+# Set default environment variable for SSH key (empty by default)
+ENV SSH_PUBLIC_KEY=""
+
+# Expose port 8888 for TensorBoard and Jupyter Notebook and port 22 for SSH
+EXPOSE 8888 22
+
+# Ensure SSH key is added when the container starts
+CMD ["/bin/bash", "-c", "\
+    mkdir -p /root/.ssh && \
+    chmod 700 /root/.ssh && \
+    echo \"$SSH_PUBLIC_KEY\" > /root/.ssh/authorized_keys && \
+    chmod 600 /root/.ssh/authorized_keys && \
+    mkdir -p /run/sshd && \
+    /usr/sbin/sshd -D"]

From ed04e004cda475461ab23dd1c9ccec59a5671037 Mon Sep 17 00:00:00 2001
From: Anton Zamyatin <e12223389@student.tuwien.ac.at>
Date: Mon, 3 Mar 2025 17:43:39 +0100
Subject: [PATCH 05/16] :whale: Reformat Dockerfile

---
 Dockerfile | 15 ++++++---------
 1 file changed, 6 insertions(+), 9 deletions(-)

diff --git a/Dockerfile b/Dockerfile
index 4ae9edca..b7b86295 100644
--- a/Dockerfile
+++ b/Dockerfile
@@ -5,7 +5,7 @@ ENV DEBIAN_FRONTEND=noninteractive \
     LC_ALL=C.UTF-8 \
     LANG=C.UTF-8
 
-# Install Git and OpenSSH Server (PyTorch's base image alraedy includes Conda and Pip)
+# Install Git and OpenSSH Server (PyTorch's base image already includes Conda and Pip)
 RUN apt-get update && apt-get install -y \
     git \
     openssh-server \
@@ -13,7 +13,7 @@ RUN apt-get update && apt-get install -y \
 
 WORKDIR /workspace
 
-# Copy README.md and torch_uncertainy module (required by pyproject.toml, othwise flit build will fail)
+# Copy README.md and torch_uncertainty module (required by pyproject.toml, otherwise flit build will fail)
 COPY README.md /workspace/
 COPY torch_uncertainty /workspace/torch_uncertainty
 
@@ -26,7 +26,7 @@ RUN pip install --no-cache-dir ".[all]"
 # Always activate Conda when opening a new terminal
 RUN echo "source /opt/conda/bin/activate" >> /root/.bashrc
 
-# Customize the Bash prompt
+# Customize shell prompt
 RUN echo 'force_color_prompt=yes' >> /root/.bashrc && \
     echo 'PS1="\[\033[01;34m\]\W\[\033[00m\]\$ "' >> /root/.bashrc && \
     echo 'if [ -x /usr/bin/dircolors ]; then' >> /root/.bashrc && \
@@ -38,22 +38,19 @@ RUN echo 'force_color_prompt=yes' >> /root/.bashrc && \
     echo '    cd /workspace' >> /root/.bashrc && \
     echo 'fi' >> /root/.bashrc
 
-# Allow SSH login as root without a password (key-based authentication only)
+# Configure SSH server
 RUN echo "PermitRootLogin yes" >> /etc/ssh/sshd_config && \
     echo "PubkeyAuthentication yes" >> /etc/ssh/sshd_config && \
     echo "AuthorizedKeysFile .ssh/authorized_keys" >> /etc/ssh/sshd_config
 
-# Set default environment variable for SSH key (empty by default)
-ENV SSH_PUBLIC_KEY=""
-
 # Expose port 8888 for TensorBoard and Jupyter Notebook and port 22 for SSH
 EXPOSE 8888 22
 
-# Ensure SSH key is added when the container starts
+# Ensure public key for RunPod-Auth is added when the container starts
 CMD ["/bin/bash", "-c", "\
     mkdir -p /root/.ssh && \
     chmod 700 /root/.ssh && \
-    echo \"$SSH_PUBLIC_KEY\" > /root/.ssh/authorized_keys && \
+    echo \"$RUNPOD_SSH_PUBLIC_KEY\" > /root/.ssh/authorized_keys && \
     chmod 600 /root/.ssh/authorized_keys && \
     mkdir -p /run/sshd && \
     /usr/sbin/sshd -D"]

From 5d731cd722a38f291351ef04372fc03f4ba49ea7 Mon Sep 17 00:00:00 2001
From: Anton Zamyatin <e12223389@student.tuwien.ac.at>
Date: Mon, 3 Mar 2025 22:18:04 +0100
Subject: [PATCH 06/16] :whale: Install OpenGL and run SSH setup for
 GitHub-Auth

---
 Dockerfile | 12 +++++++++++-
 1 file changed, 11 insertions(+), 1 deletion(-)

diff --git a/Dockerfile b/Dockerfile
index b7b86295..8c88e414 100644
--- a/Dockerfile
+++ b/Dockerfile
@@ -5,10 +5,11 @@ ENV DEBIAN_FRONTEND=noninteractive \
     LC_ALL=C.UTF-8 \
     LANG=C.UTF-8
 
-# Install Git and OpenSSH Server (PyTorch's base image already includes Conda and Pip)
+# Install Git, OpenSSH Server, and OpenGL (PyTorch's base image already includes Conda and Pip)
 RUN apt-get update && apt-get install -y \
     git \
     openssh-server \
+    libgl1 \
     && rm -rf /var/lib/apt/lists/*
 
 WORKDIR /workspace
@@ -52,5 +53,14 @@ CMD ["/bin/bash", "-c", "\
     chmod 700 /root/.ssh && \
     echo \"$RUNPOD_SSH_PUBLIC_KEY\" > /root/.ssh/authorized_keys && \
     chmod 600 /root/.ssh/authorized_keys && \
+    echo \"$GITHUB_SSH_PRIVATE_KEY\" > /root/.ssh/github_rsa && \
+    chmod 600 /root/.ssh/github_rsa && \
+    echo 'Host github.com' > /root/.ssh/config && \
+    echo '  User git' >> /root/.ssh/config && \
+    echo '  IdentityFile /root/.ssh/github_rsa' >> /root/.ssh/config && \
+    chmod 600 /root/.ssh/config && \
+    eval $(ssh-agent -s) && ssh-add /root/.ssh/github_rsa && \
+    ssh-keyscan github.com >> /root/.ssh/known_hosts && \
+    git clone git@github.com:$GITHUB_USER/torch-uncertainty.git /workspace && \
     mkdir -p /run/sshd && \
     /usr/sbin/sshd -D"]

From 025f84dadedb4a6bc99a16c0260872cadae50b78 Mon Sep 17 00:00:00 2001
From: Anton Zamyatin <e12223389@student.tuwien.ac.at>
Date: Mon, 3 Mar 2025 23:46:46 +0100
Subject: [PATCH 07/16] :whale: Ensure setup only runs once when container is
 started for ther first time

---
 Dockerfile | 46 +++++++++++++++++++++++++++++-----------------
 1 file changed, 29 insertions(+), 17 deletions(-)

diff --git a/Dockerfile b/Dockerfile
index 8c88e414..cd288ee3 100644
--- a/Dockerfile
+++ b/Dockerfile
@@ -27,15 +27,18 @@ RUN pip install --no-cache-dir ".[all]"
 # Always activate Conda when opening a new terminal
 RUN echo "source /opt/conda/bin/activate" >> /root/.bashrc
 
-# Customize shell prompt
+# Customize shell prompt (optional)
 RUN echo 'force_color_prompt=yes' >> /root/.bashrc && \
+    # Blue working directory, no username, and no hostname, with $ at the end
     echo 'PS1="\[\033[01;34m\]\W\[\033[00m\]\$ "' >> /root/.bashrc && \
+    # Colorize ls, grep, fgrep, and egrep
     echo 'if [ -x /usr/bin/dircolors ]; then' >> /root/.bashrc && \
     echo '    test -r ~/.dircolors && eval "$(dircolors -b ~/.dircolors)" || eval "$(dircolors -b)"' >> /root/.bashrc && \
     echo '    alias ls="ls --color=auto"' >> /root/.bashrc && \
     echo '    alias grep="grep --color=auto"' >> /root/.bashrc && \
     echo '    alias fgrep="fgrep --color=auto"' >> /root/.bashrc && \
     echo '    alias egrep="egrep --color=auto"' >> /root/.bashrc && \
+    # Automatically change to workspace directory when opening a new terminal
     echo '    cd /workspace' >> /root/.bashrc && \
     echo 'fi' >> /root/.bashrc
 
@@ -47,20 +50,29 @@ RUN echo "PermitRootLogin yes" >> /etc/ssh/sshd_config && \
 # Expose port 8888 for TensorBoard and Jupyter Notebook and port 22 for SSH
 EXPOSE 8888 22
 
-# Ensure public key for RunPod-Auth is added when the container starts
 CMD ["/bin/bash", "-c", "\
-    mkdir -p /root/.ssh && \
-    chmod 700 /root/.ssh && \
-    echo \"$RUNPOD_SSH_PUBLIC_KEY\" > /root/.ssh/authorized_keys && \
-    chmod 600 /root/.ssh/authorized_keys && \
-    echo \"$GITHUB_SSH_PRIVATE_KEY\" > /root/.ssh/github_rsa && \
-    chmod 600 /root/.ssh/github_rsa && \
-    echo 'Host github.com' > /root/.ssh/config && \
-    echo '  User git' >> /root/.ssh/config && \
-    echo '  IdentityFile /root/.ssh/github_rsa' >> /root/.ssh/config && \
-    chmod 600 /root/.ssh/config && \
-    eval $(ssh-agent -s) && ssh-add /root/.ssh/github_rsa && \
-    ssh-keyscan github.com >> /root/.ssh/known_hosts && \
-    git clone git@github.com:$GITHUB_USER/torch-uncertainty.git /workspace && \
-    mkdir -p /run/sshd && \
-    /usr/sbin/sshd -D"]
+    # Ensure first-time setup only runs once
+    if [ ! -f /workspace/.setup_done ]; then \
+        echo 'Running first-time setup...'; \
+        # Add public key for RunPod-Auth and private key for GitHub-Auth
+        mkdir -p /root/.ssh && chmod 700 /root/.ssh; \
+        echo \"$RUNPOD_SSH_PUBLIC_KEY\" > /root/.ssh/authorized_keys && chmod 600 /root/.ssh/authorized_keys; \
+        echo \"$GITHUB_SSH_PRIVATE_KEY\" > /root/.ssh/github_rsa && chmod 600 /root/.ssh/github_rsa; \
+        # Add GitHub credentials to SSh config
+        echo 'Host github.com' > /root/.ssh/config; \
+        echo '  User git' >> /root/.ssh/config; \
+        echo '  IdentityFile /root/.ssh/github_rsa' >> /root/.ssh/config; \
+        chmod 600 /root/.ssh/config; \
+        # Add GitHub to known hosts
+        ssh-keyscan github.com >> /root/.ssh/known_hosts; \
+        # Clone GitHub repo if not already cloned
+        if [ ! -d \"/workspace/.git\" ]; then git clone git@github.com:$GITHUB_USER/torch-uncertainty.git /workspace; fi; \
+        # Mark first-time setup as done
+        touch /workspace/.setup_done; \
+    else \
+        echo 'Skipping first-time setup, already done.'; \
+    fi; \
+    # Always start SSH agent and add key every time the container starts
+    eval $(ssh-agent -s) && ssh-add /root/.ssh/github_rsa; \
+    # Start SSH server
+    mkdir -p /run/sshd && /usr/sbin/sshd -D"]

From 4106501a50235657a964975d23355fabd4a70a7c Mon Sep 17 00:00:00 2001
From: Anton Zamyatin <e12223389@student.tuwien.ac.at>
Date: Tue, 4 Mar 2025 09:51:41 +0100
Subject: [PATCH 08/16] :whale: Optimize setup steps and make custom prompt
 optional via env variable

---
 .gitignore |  1 +
 Dockerfile | 70 +++++++++++++++++++++++++++++++++++++++++-------------
 2 files changed, 55 insertions(+), 16 deletions(-)

diff --git a/.gitignore b/.gitignore
index 9ad40cf4..fcd341af 100644
--- a/.gitignore
+++ b/.gitignore
@@ -41,6 +41,7 @@ share/python-wheels/
 .installed.cfg
 *.egg
 MANIFEST
+.setup_done
 
 # PyInstaller
 #  Usually these files are written by a python script from a template
diff --git a/Dockerfile b/Dockerfile
index cd288ee3..8618c582 100644
--- a/Dockerfile
+++ b/Dockerfile
@@ -28,7 +28,8 @@ RUN pip install --no-cache-dir ".[all]"
 RUN echo "source /opt/conda/bin/activate" >> /root/.bashrc
 
 # Customize shell prompt (optional)
-RUN echo 'force_color_prompt=yes' >> /root/.bashrc && \
+RUN if [ ! -z "$USE_COMPACT_SHELL_PROMPT" ] && [ "$USE_COMPACT_SHELL_PROMPT" = "true" ]; then \
+    echo 'force_color_prompt=yes' >> /root/.bashrc && \
     # Blue working directory, no username, and no hostname, with $ at the end
     echo 'PS1="\[\033[01;34m\]\W\[\033[00m\]\$ "' >> /root/.bashrc && \
     # Colorize ls, grep, fgrep, and egrep
@@ -40,7 +41,8 @@ RUN echo 'force_color_prompt=yes' >> /root/.bashrc && \
     echo '    alias egrep="egrep --color=auto"' >> /root/.bashrc && \
     # Automatically change to workspace directory when opening a new terminal
     echo '    cd /workspace' >> /root/.bashrc && \
-    echo 'fi' >> /root/.bashrc
+    echo 'fi' >> /root/.bashrc \
+    fi;
 
 # Configure SSH server
 RUN echo "PermitRootLogin yes" >> /etc/ssh/sshd_config && \
@@ -50,29 +52,65 @@ RUN echo "PermitRootLogin yes" >> /etc/ssh/sshd_config && \
 # Expose port 8888 for TensorBoard and Jupyter Notebook and port 22 for SSH
 EXPOSE 8888 22
 
+# Entrypoint script (runs every time the container starts)
 CMD ["/bin/bash", "-c", "\
-    # Ensure first-time setup only runs once
-    if [ ! -f /workspace/.setup_done ]; then \
-        echo 'Running first-time setup...'; \
-        # Add public key for RunPod-Auth and private key for GitHub-Auth
+    # Create SSH directory and set permissions if not present \
+        if [ ! -d /root/.ssh ]; then \
         mkdir -p /root/.ssh && chmod 700 /root/.ssh; \
+    fi; \
+    # Add public key for RunPod-Auth if not present \
+    if [ -z \"$RUNPOD_SSH_PUBLIC_KEY\" ]; then \
+        echo 'Please set the RUNPOD_SSH_PUBLIC_KEY environment variable.'; \
+        exit 1; \
+    fi; \
+    if [ ! -f /root/.ssh/authorized_keys ] || ! grep -q \"$RUNPOD_SSH_PUBLIC_KEY\" /root/.ssh/authorized_keys; then \
         echo \"$RUNPOD_SSH_PUBLIC_KEY\" > /root/.ssh/authorized_keys && chmod 600 /root/.ssh/authorized_keys; \
+    fi; \
+    # Add private key for GitHub-Auth if not present \
+    if [ -z \"$GITHUB_SSH_PRIVATE_KEY\" ]; then \
+        echo 'Please set the GITHUB_SSH_PRIVATE_KEY environment variable.'; \
+        exit 1; \
+    fi; \
+    if [ ! -f /root/.ssh/github_rsa ]; then \
         echo \"$GITHUB_SSH_PRIVATE_KEY\" > /root/.ssh/github_rsa && chmod 600 /root/.ssh/github_rsa; \
-        # Add GitHub credentials to SSh config
+    fi; \
+    # Add GitHub credentials to SSH config if not present \
+    if [ ! -f /root/.ssh/config ] || ! grep -q 'Host github.com' /root/.ssh/config; then \
         echo 'Host github.com' > /root/.ssh/config; \
         echo '  User git' >> /root/.ssh/config; \
         echo '  IdentityFile /root/.ssh/github_rsa' >> /root/.ssh/config; \
         chmod 600 /root/.ssh/config; \
-        # Add GitHub to known hosts
-        ssh-keyscan github.com >> /root/.ssh/known_hosts; \
-        # Clone GitHub repo if not already cloned
-        if [ ! -d \"/workspace/.git\" ]; then git clone git@github.com:$GITHUB_USER/torch-uncertainty.git /workspace; fi; \
-        # Mark first-time setup as done
+    fi; \
+    # Add GitHub to known hosts if not already added \
+    ssh-keygen -F github.com > /dev/null 2>&1 || ssh-keyscan github.com >> /root/.ssh/known_hosts; \
+    # Start SSH agent if not running and add GitHub private key \
+    if ! pgrep -x \"ssh-agent\" > /dev/null; then \
+        eval $(ssh-agent -s); \
+    fi; \
+    ssh-add -l | grep github_rsa > /dev/null || ssh-add /root/.ssh/github_rsa; \
+    # Ensure first-time setup only runs once \
+    if [ ! -f /workspace/.setup_done ]; then \
+        echo 'Running first-time setup...'; \
+        # Clone GitHub repo if not already cloned \
+        if [ -z \"$GITHUB_USER\" ]; then \
+            echo 'Please set the GITHUB_USER environment variable.'; \
+            exit 1; \
+        fi; \
+        if [ ! -d \"/workspace/.git\" ]; then \
+            git clone git@github.com:$GITHUB_USER/torch-uncertainty.git /workspace; \
+        fi; \
+        # Set Git user name and email if provided \
+        if [ ! -z \"$GIT_USER_NAME\" ]; then \
+            git config --global user.name \"$GIT_USER_NAME\"; \
+        fi; \
+        if [ ! -z \"$GIT_USER_EMAIL\" ]; then \
+            git config --global user.email \"$GIT_USER_EMAIL\"; \
+        fi; \
+        # Mark first-time setup as done \
         touch /workspace/.setup_done; \
     else \
         echo 'Skipping first-time setup, already done.'; \
     fi; \
-    # Always start SSH agent and add key every time the container starts
-    eval $(ssh-agent -s) && ssh-add /root/.ssh/github_rsa; \
-    # Start SSH server
-    mkdir -p /run/sshd && /usr/sbin/sshd -D"]
+    # Start SSH server \
+    mkdir -p /run/sshd && chmod 755 /run/sshd; \
+    /usr/sbin/sshd -D"]

From c8cd50b5f019f86a4411e4c85e3dd8a5cb78802e Mon Sep 17 00:00:00 2001
From: Anton Zamyatin <e12223389@student.tuwien.ac.at>
Date: Tue, 4 Mar 2025 13:32:24 +0100
Subject: [PATCH 09/16] :whale: Optimize build time and run shell prompt in
 initial setup

---
 Dockerfile | 49 ++++++++++++++++++++++++-------------------------
 1 file changed, 24 insertions(+), 25 deletions(-)

diff --git a/Dockerfile b/Dockerfile
index 8618c582..6bfd1f74 100644
--- a/Dockerfile
+++ b/Dockerfile
@@ -14,9 +14,8 @@ RUN apt-get update && apt-get install -y \
 
 WORKDIR /workspace
 
-# Copy README.md and torch_uncertainty module (required by pyproject.toml, otherwise flit build will fail)
-COPY README.md /workspace/
-COPY torch_uncertainty /workspace/torch_uncertainty
+# Create an empty README.md file and an empty torch_uncertainty module to satisfy flit
+RUN touch README.md && mkdir -p torch_uncertainty && touch torch_uncertainty/__init__.py
 
 # Copy dependency file
 COPY pyproject.toml /workspace/
@@ -27,23 +26,6 @@ RUN pip install --no-cache-dir ".[all]"
 # Always activate Conda when opening a new terminal
 RUN echo "source /opt/conda/bin/activate" >> /root/.bashrc
 
-# Customize shell prompt (optional)
-RUN if [ ! -z "$USE_COMPACT_SHELL_PROMPT" ] && [ "$USE_COMPACT_SHELL_PROMPT" = "true" ]; then \
-    echo 'force_color_prompt=yes' >> /root/.bashrc && \
-    # Blue working directory, no username, and no hostname, with $ at the end
-    echo 'PS1="\[\033[01;34m\]\W\[\033[00m\]\$ "' >> /root/.bashrc && \
-    # Colorize ls, grep, fgrep, and egrep
-    echo 'if [ -x /usr/bin/dircolors ]; then' >> /root/.bashrc && \
-    echo '    test -r ~/.dircolors && eval "$(dircolors -b ~/.dircolors)" || eval "$(dircolors -b)"' >> /root/.bashrc && \
-    echo '    alias ls="ls --color=auto"' >> /root/.bashrc && \
-    echo '    alias grep="grep --color=auto"' >> /root/.bashrc && \
-    echo '    alias fgrep="fgrep --color=auto"' >> /root/.bashrc && \
-    echo '    alias egrep="egrep --color=auto"' >> /root/.bashrc && \
-    # Automatically change to workspace directory when opening a new terminal
-    echo '    cd /workspace' >> /root/.bashrc && \
-    echo 'fi' >> /root/.bashrc \
-    fi;
-
 # Configure SSH server
 RUN echo "PermitRootLogin yes" >> /etc/ssh/sshd_config && \
     echo "PubkeyAuthentication yes" >> /etc/ssh/sshd_config && \
@@ -58,14 +40,16 @@ CMD ["/bin/bash", "-c", "\
         if [ ! -d /root/.ssh ]; then \
         mkdir -p /root/.ssh && chmod 700 /root/.ssh; \
     fi; \
-    # Add public key for RunPod-Auth if not present \
-    if [ -z \"$RUNPOD_SSH_PUBLIC_KEY\" ]; then \
-        echo 'Please set the RUNPOD_SSH_PUBLIC_KEY environment variable.'; \
+    \
+    # Add public key for VM-Auth if not present \
+    if [ -z \"$VM_SSH_PUBLIC_KEY\" ]; then \
+        echo 'Please set the VM_SSH_PUBLIC_KEY environment variable.'; \
         exit 1; \
     fi; \
-    if [ ! -f /root/.ssh/authorized_keys ] || ! grep -q \"$RUNPOD_SSH_PUBLIC_KEY\" /root/.ssh/authorized_keys; then \
-        echo \"$RUNPOD_SSH_PUBLIC_KEY\" > /root/.ssh/authorized_keys && chmod 600 /root/.ssh/authorized_keys; \
+    if [ ! -f /root/.ssh/authorized_keys ] || ! grep -q \"$VM_SSH_PUBLIC_KEY\" /root/.ssh/authorized_keys; then \
+        echo \"$VM_SSH_PUBLIC_KEY\" > /root/.ssh/authorized_keys && chmod 600 /root/.ssh/authorized_keys; \
     fi; \
+    \
     # Add private key for GitHub-Auth if not present \
     if [ -z \"$GITHUB_SSH_PRIVATE_KEY\" ]; then \
         echo 'Please set the GITHUB_SSH_PRIVATE_KEY environment variable.'; \
@@ -88,6 +72,7 @@ CMD ["/bin/bash", "-c", "\
         eval $(ssh-agent -s); \
     fi; \
     ssh-add -l | grep github_rsa > /dev/null || ssh-add /root/.ssh/github_rsa; \
+    \
     # Ensure first-time setup only runs once \
     if [ ! -f /workspace/.setup_done ]; then \
         echo 'Running first-time setup...'; \
@@ -111,6 +96,20 @@ CMD ["/bin/bash", "-c", "\
     else \
         echo 'Skipping first-time setup, already done.'; \
     fi; \
+    \
+    # Apply shell prompt customization \
+    if [ ! -z \"$USE_COMPACT_SHELL_PROMPT\" ]; then \
+        echo 'force_color_prompt=yes' >> /root/.bashrc; \
+        echo 'PS1=\"\[\033[01;34m\]\W\[\033[00m\]\$ \"' >> /root/.bashrc; \
+        echo 'if [ -x /usr/bin/dircolors ]; then' >> /root/.bashrc; \
+        echo '    test -r ~/.dircolors && eval \"$(dircolors -b ~/.dircolors)\" || eval \"$(dircolors -b)\"' >> /root/.bashrc; \
+        echo '    alias ls=\"ls --color=auto\"' >> /root/.bashrc; \
+        echo '    alias grep=\"grep --color=auto\"' >> /root/.bashrc; \
+        echo '    alias fgrep=\"fgrep --color=auto\"' >> /root/.bashrc; \
+        echo '    alias egrep=\"egrep --color=auto\"' >> /root/.bashrc; \
+        echo 'fi' >> /root/.bashrc; \
+    fi; \
+    \
     # Start SSH server \
     mkdir -p /run/sshd && chmod 755 /run/sshd; \
     /usr/sbin/sshd -D"]

From a49bd3ec0477f2cc3b2adfa6fad119e8d348ab3b Mon Sep 17 00:00:00 2001
From: Anton Zamyatin <e12223389@student.tuwien.ac.at>
Date: Tue, 4 Mar 2025 13:48:28 +0100
Subject: [PATCH 10/16] :books: Add documentation for using Docker image

---
 DOCKER.md | 69 +++++++++++++++++++++++++++++++++++++++++++++++++++++++
 README.md |  6 +++++
 2 files changed, 75 insertions(+)
 create mode 100644 DOCKER.md

diff --git a/DOCKER.md b/DOCKER.md
new file mode 100644
index 00000000..1b53beb4
--- /dev/null
+++ b/DOCKER.md
@@ -0,0 +1,69 @@
+# :whale: Docker image for contributors
+### Pre-built Docker image
+1. To pull the pre-built image from Docker Hub, simply run:
+    ```bash
+    docker pull docker.io/tonyzamyatin/torch-uncertainty:latest
+    ```
+
+    This image includes:
+    - PyTorch with CUDA support
+    - OpenGL (for visualization tasks)
+    - Git, OpenSSH, and all Python dependencies
+
+    Checkout the [registry on Docker Hub](https://hub.docker.com/repository/docker/tonyzamyatin/torch-uncertainty/general) for all available images.
+
+2. To start a container using this image, set up the necessary environment variables and run:
+    ```bash
+    docker run --rm -it --gpus all -p 8888:8888 -p 22:22 \
+        -e VM_SSH_PUBLIC_KEY="your-public-key" \
+        -e GITHUB_SSH_PRIVATE_KEY="your-github-key" \
+        -e GITHUB_USER="your-github-username" \
+        -e GIT_USER_EMAIL="your-git-email" \
+        -e GIT_USER_NAME="your-git-name" \
+        docker.io/tonyzamyatin/torch-uncertainty
+    ```
+
+    Optionally, you can also set `-e USER_COMPACT_SHELL_PROMPT="true"`  
+    to make the VM's shell prompts compact and colorized.
+
+    **Note:** Some cloud providers offer templates, in which you can preconfigure  
+    in advance which Docker image to pull and which environment variables to set.  
+    In this case, the provider will pull the image, set all environment variables,  
+    and start the container for you.
+
+3. Once your cloud provider has deployed the VM, it will display the host address and SSH port.  
+    You can connect to the container via SSH using:
+    ```bash
+    ssh -i /path/to/private_key root@<VM_HOST> -p <VM_PORT>
+    ```
+
+    Replace `<VM_HOST>` and `<VM_PORT>` with the values provided by your cloud provider,  
+    and `/path/to/private_key` with the private key that corresponds to `VM_SSH_PUBLIC_KEY`.
+
+4. The container exposes port `8888` in case you want to run Jupyter Notebooks or TensorBoard.
+
+    **Note:** The `/workspace` directory is mounted from your local machine or cloud storage,  
+    so changes persist across container restarts.  
+    If using a cloud provider, ensure your network volume is correctly attached to avoid losing data.
+
+### Modifying and publishing custom Docker image
+If you want to make changes to the Dockerfile, follow these steps:
+1. Edit the Dockerfile to fit your needs.
+
+2. Build the modified image:
+    ```
+    docker build -t my-custom-image .
+    ```
+
+3. Push to a Docker registry (if you want to use it on another VM):
+    ```
+    docker tag my-custom-image mydockerhubuser/my-custom-image:tag
+    docker push mydockerhubuser/my-custom-image:tag
+    ```
+    
+4. Pull the custom image onto your VM:
+    ```
+    docker pull mydockerhubuser/my-custom-image
+    ```
+    
+5. Run the container using the same docker run command with the new image name.
diff --git a/README.md b/README.md
index 52b0ed8e..aca06bf3 100644
--- a/README.md
+++ b/README.md
@@ -95,3 +95,9 @@ The following projects use TorchUncertainty:
 - *A Symmetry-Aware Exploration of Bayesian Neural Network Posteriors* - [ICLR 2024](https://arxiv.org/abs/2310.08287)
 
 **If you are using TorchUncertainty in your project, please let us know, we will add your project to this list!**
+
+## :whale: Docker image for contributors
+For contributors who want to run experiments on cloud GPU instances, we provide a pre-built Docker image that includes all necessary dependencies and configurations and the Dockerfile for building your custom Docker images. 
+This allows you to quickly launch an experiment-ready container with minimal setup.
+
+Please refer to [DOCKER.md](DOCKER.md) for further details.

From bb79e0bf1071e73469ddde403f2f774cef023542 Mon Sep 17 00:00:00 2001
From: Anton Zamyatin <e12223389@student.tuwien.ac.at>
Date: Tue, 4 Mar 2025 18:39:24 +0100
Subject: [PATCH 11/16] :whale: Install dependencies in editable mode

---
 Dockerfile | 5 ++---
 1 file changed, 2 insertions(+), 3 deletions(-)

diff --git a/Dockerfile b/Dockerfile
index 6bfd1f74..255291d1 100644
--- a/Dockerfile
+++ b/Dockerfile
@@ -20,8 +20,8 @@ RUN touch README.md && mkdir -p torch_uncertainty && touch torch_uncertainty/__i
 # Copy dependency file
 COPY pyproject.toml /workspace/
 
-# Install dependencies
-RUN pip install --no-cache-dir ".[all]"
+# Install dependencies (in editable mode!)
+RUN pip install --no-cache-dir "-e .[all]"
 
 # Always activate Conda when opening a new terminal
 RUN echo "source /opt/conda/bin/activate" >> /root/.bashrc
@@ -109,7 +109,6 @@ CMD ["/bin/bash", "-c", "\
         echo '    alias egrep=\"egrep --color=auto\"' >> /root/.bashrc; \
         echo 'fi' >> /root/.bashrc; \
     fi; \
-    \
     # Start SSH server \
     mkdir -p /run/sshd && chmod 755 /run/sshd; \
     /usr/sbin/sshd -D"]

From eaa09d641b775810ab66bfb413b068823e218279 Mon Sep 17 00:00:00 2001
From: Anton Zamyatin <e12223389@student.tuwien.ac.at>
Date: Tue, 4 Mar 2025 22:16:24 +0100
Subject: [PATCH 12/16] :bug: Fix pip install in editable mode

---
 Dockerfile | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/Dockerfile b/Dockerfile
index 255291d1..db632e3a 100644
--- a/Dockerfile
+++ b/Dockerfile
@@ -21,7 +21,7 @@ RUN touch README.md && mkdir -p torch_uncertainty && touch torch_uncertainty/__i
 COPY pyproject.toml /workspace/
 
 # Install dependencies (in editable mode!)
-RUN pip install --no-cache-dir "-e .[all]"
+RUN pip install --no-cache-dir -e ".[all]"
 
 # Always activate Conda when opening a new terminal
 RUN echo "source /opt/conda/bin/activate" >> /root/.bashrc

From ae0405ad999bbb94c68c483430740dbf073f1a7a Mon Sep 17 00:00:00 2001
From: Anton Zamyatin <e12223389@student.tuwien.ac.at>
Date: Wed, 5 Mar 2025 01:22:12 +0100
Subject: [PATCH 13/16] :hammer: Fix and format container start script and move
 it to entrypoint.sh

---
 Dockerfile                    |  87 +++--------------------
 README.md                     |   2 +-
 DOCKER.md => docker/DOCKER.md |   0
 docker/entrypoint.sh          | 127 ++++++++++++++++++++++++++++++++++
 4 files changed, 136 insertions(+), 80 deletions(-)
 rename DOCKER.md => docker/DOCKER.md (100%)
 create mode 100644 docker/entrypoint.sh

diff --git a/Dockerfile b/Dockerfile
index db632e3a..cc89a22b 100644
--- a/Dockerfile
+++ b/Dockerfile
@@ -20,8 +20,8 @@ RUN touch README.md && mkdir -p torch_uncertainty && touch torch_uncertainty/__i
 # Copy dependency file
 COPY pyproject.toml /workspace/
 
-# Install dependencies (in editable mode!)
-RUN pip install --no-cache-dir -e ".[all]"
+# Install dependencies all dependencies
+RUN pip install --no-cache-dir -e ".[dev]"
 
 # Always activate Conda when opening a new terminal
 RUN echo "source /opt/conda/bin/activate" >> /root/.bashrc
@@ -35,80 +35,9 @@ RUN echo "PermitRootLogin yes" >> /etc/ssh/sshd_config && \
 EXPOSE 8888 22
 
 # Entrypoint script (runs every time the container starts)
-CMD ["/bin/bash", "-c", "\
-    # Create SSH directory and set permissions if not present \
-        if [ ! -d /root/.ssh ]; then \
-        mkdir -p /root/.ssh && chmod 700 /root/.ssh; \
-    fi; \
-    \
-    # Add public key for VM-Auth if not present \
-    if [ -z \"$VM_SSH_PUBLIC_KEY\" ]; then \
-        echo 'Please set the VM_SSH_PUBLIC_KEY environment variable.'; \
-        exit 1; \
-    fi; \
-    if [ ! -f /root/.ssh/authorized_keys ] || ! grep -q \"$VM_SSH_PUBLIC_KEY\" /root/.ssh/authorized_keys; then \
-        echo \"$VM_SSH_PUBLIC_KEY\" > /root/.ssh/authorized_keys && chmod 600 /root/.ssh/authorized_keys; \
-    fi; \
-    \
-    # Add private key for GitHub-Auth if not present \
-    if [ -z \"$GITHUB_SSH_PRIVATE_KEY\" ]; then \
-        echo 'Please set the GITHUB_SSH_PRIVATE_KEY environment variable.'; \
-        exit 1; \
-    fi; \
-    if [ ! -f /root/.ssh/github_rsa ]; then \
-        echo \"$GITHUB_SSH_PRIVATE_KEY\" > /root/.ssh/github_rsa && chmod 600 /root/.ssh/github_rsa; \
-    fi; \
-    # Add GitHub credentials to SSH config if not present \
-    if [ ! -f /root/.ssh/config ] || ! grep -q 'Host github.com' /root/.ssh/config; then \
-        echo 'Host github.com' > /root/.ssh/config; \
-        echo '  User git' >> /root/.ssh/config; \
-        echo '  IdentityFile /root/.ssh/github_rsa' >> /root/.ssh/config; \
-        chmod 600 /root/.ssh/config; \
-    fi; \
-    # Add GitHub to known hosts if not already added \
-    ssh-keygen -F github.com > /dev/null 2>&1 || ssh-keyscan github.com >> /root/.ssh/known_hosts; \
-    # Start SSH agent if not running and add GitHub private key \
-    if ! pgrep -x \"ssh-agent\" > /dev/null; then \
-        eval $(ssh-agent -s); \
-    fi; \
-    ssh-add -l | grep github_rsa > /dev/null || ssh-add /root/.ssh/github_rsa; \
-    \
-    # Ensure first-time setup only runs once \
-    if [ ! -f /workspace/.setup_done ]; then \
-        echo 'Running first-time setup...'; \
-        # Clone GitHub repo if not already cloned \
-        if [ -z \"$GITHUB_USER\" ]; then \
-            echo 'Please set the GITHUB_USER environment variable.'; \
-            exit 1; \
-        fi; \
-        if [ ! -d \"/workspace/.git\" ]; then \
-            git clone git@github.com:$GITHUB_USER/torch-uncertainty.git /workspace; \
-        fi; \
-        # Set Git user name and email if provided \
-        if [ ! -z \"$GIT_USER_NAME\" ]; then \
-            git config --global user.name \"$GIT_USER_NAME\"; \
-        fi; \
-        if [ ! -z \"$GIT_USER_EMAIL\" ]; then \
-            git config --global user.email \"$GIT_USER_EMAIL\"; \
-        fi; \
-        # Mark first-time setup as done \
-        touch /workspace/.setup_done; \
-    else \
-        echo 'Skipping first-time setup, already done.'; \
-    fi; \
-    \
-    # Apply shell prompt customization \
-    if [ ! -z \"$USE_COMPACT_SHELL_PROMPT\" ]; then \
-        echo 'force_color_prompt=yes' >> /root/.bashrc; \
-        echo 'PS1=\"\[\033[01;34m\]\W\[\033[00m\]\$ \"' >> /root/.bashrc; \
-        echo 'if [ -x /usr/bin/dircolors ]; then' >> /root/.bashrc; \
-        echo '    test -r ~/.dircolors && eval \"$(dircolors -b ~/.dircolors)\" || eval \"$(dircolors -b)\"' >> /root/.bashrc; \
-        echo '    alias ls=\"ls --color=auto\"' >> /root/.bashrc; \
-        echo '    alias grep=\"grep --color=auto\"' >> /root/.bashrc; \
-        echo '    alias fgrep=\"fgrep --color=auto\"' >> /root/.bashrc; \
-        echo '    alias egrep=\"egrep --color=auto\"' >> /root/.bashrc; \
-        echo 'fi' >> /root/.bashrc; \
-    fi; \
-    # Start SSH server \
-    mkdir -p /run/sshd && chmod 755 /run/sshd; \
-    /usr/sbin/sshd -D"]
+COPY docker/entrypoint.sh /usr/local/bin/
+RUN chmod +x /usr/local/bin/entrypoint.sh
+ENTRYPOINT ["/usr/local/bin/entrypoint.sh"]
+
+# Note that if /workspace/ is a mounted volume, any files copied to /workspace/ during the build will be overwritten by the mounted volume
+# This is why we copy the entrypoint script to /usr/local/bin/ instead of /workspace/
diff --git a/README.md b/README.md
index aca06bf3..9130926e 100644
--- a/README.md
+++ b/README.md
@@ -100,4 +100,4 @@ The following projects use TorchUncertainty:
 For contributors who want to run experiments on cloud GPU instances, we provide a pre-built Docker image that includes all necessary dependencies and configurations and the Dockerfile for building your custom Docker images. 
 This allows you to quickly launch an experiment-ready container with minimal setup.
 
-Please refer to [DOCKER.md](DOCKER.md) for further details.
+Please refer to [DOCKER.md](docker/DOCKER.md) for further details.
diff --git a/DOCKER.md b/docker/DOCKER.md
similarity index 100%
rename from DOCKER.md
rename to docker/DOCKER.md
diff --git a/docker/entrypoint.sh b/docker/entrypoint.sh
new file mode 100644
index 00000000..ffef3e62
--- /dev/null
+++ b/docker/entrypoint.sh
@@ -0,0 +1,127 @@
+#!/bin/bash
+set -e  # Exit immediately if a command fails
+
+echo "🚀 Starting container..."
+
+# Ensure SSH directory exists and has correct permissions
+if [ ! -d /root/.ssh ]; then
+    echo "📂 Creating SSH directory..."
+    mkdir -p /root/.ssh && chmod 700 /root/.ssh
+fi
+
+# Ensure the VM's public SSH key is added for authentication
+if [ -z "$VM_SSH_PUBLIC_KEY" ]; then
+    echo "❌ Error: Please set the VM_SSH_PUBLIC_KEY environment variable."
+    exit 1
+fi
+if [ ! -f /root/.ssh/authorized_keys ] || ! grep -q "$VM_SSH_PUBLIC_KEY" /root/.ssh/authorized_keys; then
+    echo "🔑 Adding VM SSH public key..."
+    echo "$VM_SSH_PUBLIC_KEY" > /root/.ssh/authorized_keys && chmod 600 /root/.ssh/authorized_keys
+fi
+
+# Ensure GitHub SSH private key is set up for authentication
+if [ -z "$GITHUB_SSH_PRIVATE_KEY" ]; then
+    echo "❌ Error: Please set the GITHUB_SSH_PRIVATE_KEY environment variable."
+    exit 1
+fi
+if [ ! -f /root/.ssh/github_rsa ]; then
+    echo "🔐 Adding GitHub SSH private key..."
+    echo "$GITHUB_SSH_PRIVATE_KEY" > /root/.ssh/github_rsa && chmod 600 /root/.ssh/github_rsa
+fi
+
+# Configure SSH client for GitHub authentication
+if [ ! -f /root/.ssh/config ] || ! grep -q 'Host github.com' /root/.ssh/config; then
+    echo "⚙️  Configuring SSH client for GitHub authentication..."
+    cat <<EOF > /root/.ssh/config
+Host github.com
+  User git
+  IdentityFile /root/.ssh/github_rsa
+EOF
+    chmod 600 /root/.ssh/config
+fi
+
+# Add GitHub to known hosts (to avoid SSH verification prompts)
+echo "📌 Ensuring GitHub is a known host..."
+ssh-keygen -F github.com > /dev/null 2>&1 || ssh-keyscan github.com >> /root/.ssh/known_hosts
+
+# Start SSH agent and add GitHub private key (if not already added)
+if ! pgrep -x "ssh-agent" > /dev/null; then
+    echo "🕵️  Starting SSH agent..."
+    eval "$(ssh-agent -s)"
+fi
+if ssh-add -l | grep -q github_rsa; then
+    echo "✅ GitHub SSH key already added."
+else
+    echo "🔑 Adding GitHub SSH key to agent..."
+    ssh-add /root/.ssh/github_rsa
+fi
+
+# Set Git user name and email (if provided)
+if [ -n "$GIT_USER_NAME" ]; then
+    echo "👤 Setting Git username: $GIT_USER_NAME"
+    git config --global user.name "$GIT_USER_NAME"
+fi
+if [ -n "$GIT_USER_EMAIL" ]; then
+    echo "📧 Setting Git email: $GIT_USER_EMAIL"
+    git config --global user.email "$GIT_USER_EMAIL"
+fi
+
+# Ensure first-time setup runs only once
+if [ ! -f /workspace/.setup_done ]; then
+    echo "🛠️ Running first-time setup..."
+
+    # Ensure GitHub username is set
+    if [ -z "$GITHUB_USER" ]; then
+        echo "❌ Error: Please set the GITHUB_USER environment variable."
+        exit 1
+    fi
+
+    # Clone GitHub repo if not already cloned
+    if [ ! -d "/workspace/.git" ]; then
+        echo "📦 Cloning repository: $GITHUB_USER/torch-uncertainty..."
+        git clone git@github.com:$GITHUB_USER/torch-uncertainty.git /workspace
+    fi
+
+    # Mark setup as completed
+    touch /workspace/.setup_done
+    echo "✅ First-time setup complete!"
+else
+    echo "⏩ Skipping first-time setup (already done)."
+fi
+
+# Apply compact shell prompt customization (if enabled)
+if [ -n "$USE_COMPACT_SHELL_PROMPT" ]; then
+    echo "🎨 Applying compact shell prompt customization..."
+    echo 'force_color_prompt=yes' >> /root/.bashrc
+    echo 'PS1="\[\033[01;34m\]\W\[\033[00m\]\$ "' >> /root/.bashrc
+    echo 'if [ -x /usr/bin/dircolors ]; then' >> /root/.bashrc
+    echo '    test -r ~/.dircolors && eval "$(dircolors -b ~/.dircolors)" || eval "$(dircolors -b)"' >> /root/.bashrc
+    echo '    alias ls="ls --color=auto"' >> /root/.bashrc
+    echo '    alias grep="grep --color=auto"' >> /root/.bashrc
+    echo '    alias fgrep="fgrep --color=auto"' >> /root/.bashrc
+    echo '    alias egrep="egrep --color=auto"' >> /root/.bashrc
+    echo 'fi' >> /root/.bashrc
+fi
+
+# Ensure /workspace is in PYTHONPATH
+if ! echo "$PYTHONPATH" | grep -q "/workspace"; then
+    echo "📌 Adding /workspace to PYTHONPATH"
+    export PYTHONPATH="/workspace:$PYTHONPATH"
+else
+    echo "✅ PYTHONPATH is already correctly set."
+fi
+
+# Check if torch_uncertainty is installed in editable mode
+if pip show torch_uncertainty | grep -q "Editable project location: /workspace"; then
+    echo "✅ torch_uncertainty is already installed in editable mode. 🎉"
+else
+    echo "🔄 Reinstalling torch_uncertainty in editable mode..."
+    pip uninstall -y torch-uncertainty
+    pip install -e /workspace
+    echo "✅ torch_uncertainty is now installed in editable mode! 🚀"
+fi
+
+# Ensure SSH server is started
+echo "🔑 Starting SSH server..."
+mkdir -p /run/sshd && chmod 755 /run/sshd
+/usr/sbin/sshd -D

From 748f698a9bc73674d780a9dbaf3acaf5cb60187f Mon Sep 17 00:00:00 2001
From: Olivier Laurent <62881275+o-laurent@users.noreply.github.com>
Date: Fri, 7 Mar 2025 11:03:46 +0100
Subject: [PATCH 14/16] :shirt: Move Docker in the installation section

---
 README.md | 12 ++++++------
 1 file changed, 6 insertions(+), 6 deletions(-)

diff --git a/README.md b/README.md
index 9130926e..e339c989 100644
--- a/README.md
+++ b/README.md
@@ -43,6 +43,11 @@ pip install torch-uncertainty
 
 The installation procedure for contributors is different: have a look at the [contribution page](https://torch-uncertainty.github.io/contributing.html).
 
+### :whale: Docker image for contributors
+
+For contributors who want to run experiments on cloud GPU instances, we provide a pre-built Docker image that includes all necessary dependencies and configurations and the Dockerfile for building your custom Docker images. 
+This allows you to quickly launch an experiment-ready container with minimal setup. Please refer to [DOCKER.md](docker/DOCKER.md) for further details.
+
 ## :racehorse: Quickstart
 
 We make a quickstart available at [torch-uncertainty.github.io/quickstart](https://torch-uncertainty.github.io/quickstart.html).
@@ -94,10 +99,5 @@ The following projects use TorchUncertainty:
 
 - *A Symmetry-Aware Exploration of Bayesian Neural Network Posteriors* - [ICLR 2024](https://arxiv.org/abs/2310.08287)
 
-**If you are using TorchUncertainty in your project, please let us know, we will add your project to this list!**
-
-## :whale: Docker image for contributors
-For contributors who want to run experiments on cloud GPU instances, we provide a pre-built Docker image that includes all necessary dependencies and configurations and the Dockerfile for building your custom Docker images. 
-This allows you to quickly launch an experiment-ready container with minimal setup.
+**If you are using TorchUncertainty in your project, please let us know, and we will add your project to this list!**
 
-Please refer to [DOCKER.md](docker/DOCKER.md) for further details.

From eda886b3aaff9880c4bbef965129c5233f39a3e7 Mon Sep 17 00:00:00 2001
From: Olivier Laurent <62881275+o-laurent@users.noreply.github.com>
Date: Fri, 7 Mar 2025 14:43:28 +0100
Subject: [PATCH 15/16] :whale: Update Dockerfile

---
 Dockerfile | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/Dockerfile b/Dockerfile
index cc89a22b..1fd3f3d8 100644
--- a/Dockerfile
+++ b/Dockerfile
@@ -21,7 +21,7 @@ RUN touch README.md && mkdir -p torch_uncertainty && touch torch_uncertainty/__i
 COPY pyproject.toml /workspace/
 
 # Install dependencies all dependencies
-RUN pip install --no-cache-dir -e ".[dev]"
+RUN pip install --no-cache-dir -e ".[all]"
 
 # Always activate Conda when opening a new terminal
 RUN echo "source /opt/conda/bin/activate" >> /root/.bashrc

From 22428050954ee7a7dcd0d1d8d631661419b3b275 Mon Sep 17 00:00:00 2001
From: Anton Zamyatin <e12223389@student.tuwien.ac.at>
Date: Fri, 7 Mar 2025 17:40:00 +0100
Subject: [PATCH 16/16] :whale: Install pre-commit hooks on container start

---
 README.md            | 1 -
 docker/entrypoint.sh | 4 ++++
 2 files changed, 4 insertions(+), 1 deletion(-)

diff --git a/README.md b/README.md
index e339c989..e0f735a4 100644
--- a/README.md
+++ b/README.md
@@ -100,4 +100,3 @@ The following projects use TorchUncertainty:
 - *A Symmetry-Aware Exploration of Bayesian Neural Network Posteriors* - [ICLR 2024](https://arxiv.org/abs/2310.08287)
 
 **If you are using TorchUncertainty in your project, please let us know, and we will add your project to this list!**
-
diff --git a/docker/entrypoint.sh b/docker/entrypoint.sh
index ffef3e62..1c1af5e5 100644
--- a/docker/entrypoint.sh
+++ b/docker/entrypoint.sh
@@ -121,6 +121,10 @@ else
     echo "✅ torch_uncertainty is now installed in editable mode! 🚀"
 fi
 
+# Activate pre-commit hooks (if enabled)
+echo "🔗 Activating pre-commit hooks..."
+pre-commit install
+
 # Ensure SSH server is started
 echo "🔑 Starting SSH server..."
 mkdir -p /run/sshd && chmod 755 /run/sshd