diff --git a/.github/workflows/tests-docker.yml b/.github/workflows/tests-docker.yml
index 988244335f..3d3439b83f 100644
--- a/.github/workflows/tests-docker.yml
+++ b/.github/workflows/tests-docker.yml
@@ -36,7 +36,7 @@ jobs:
           PYTORCH_VERSION="${{ matrix.pytorch }}"
           # Build the Docker image
           docker build . \
-            --file ./docker/Dockerfile \
+            --file ./docker/Dockerfile-tests \
             --build-arg BASE_TAG=$BASE_TAG \
             --build-arg CUDA=$CUDA \
             --build-arg GITHUB_REF=$GITHUB_REF \
diff --git a/docker/Dockerfile b/docker/Dockerfile
index 267f8418fd..f8e0528562 100644
--- a/docker/Dockerfile
+++ b/docker/Dockerfile
@@ -6,7 +6,6 @@ ARG AXOLOTL_EXTRAS=""
 ARG CUDA="118"
 ENV BNB_CUDA_VERSION=$CUDA
 ARG PYTORCH_VERSION="2.0.1"
-ARG GITHUB_REF="main"
 
 ENV PYTORCH_VERSION=$PYTORCH_VERSION
 
@@ -15,7 +14,7 @@ RUN apt-get update && \
 
 WORKDIR /workspace
 
-RUN git clone --depth=1 -b $GITHUB_REF https://github.com/OpenAccess-AI-Collective/axolotl.git
+RUN git clone --depth=1 https://github.com/OpenAccess-AI-Collective/axolotl.git
 
 WORKDIR /workspace/axolotl