From a8e5b9d8dbe829f1c11165406281243222a78812 Mon Sep 17 00:00:00 2001
From: Irene Dea <deaairene@gmail.com>
Date: Tue, 30 Jan 2024 23:57:11 -0800
Subject: [PATCH] fix branch name

---
 .github/workflows/docker.yaml | 3 ++-
 Dockerfile                    | 3 ++-
 2 files changed, 4 insertions(+), 2 deletions(-)

diff --git a/.github/workflows/docker.yaml b/.github/workflows/docker.yaml
index 013f9b3682..e0b2de4335 100644
--- a/.github/workflows/docker.yaml
+++ b/.github/workflows/docker.yaml
@@ -72,12 +72,13 @@ jobs:
 
     - name: Build and Push the Docker Image
       uses: docker/build-push-action@v3
-      with:
+      with:       
         context: .
         tags: ${{ env.IMAGE_TAG }}
         push: true
         cache-from: type=registry,ref=${{ env.IMAGE_CACHE }}
         cache-to: type=registry,ref=${{ env.IMAGE_CACHE }},mode=max
         build-args: |
+          BRANCH_NAME=${{ github.head_ref || github.ref_name }} 
           BASE_IMAGE=${{ matrix.base_image }}
           DEP_GROUPS=${{ matrix.dep_groups }}
diff --git a/Dockerfile b/Dockerfile
index 430a2cae03..90d255a33d 100644
--- a/Dockerfile
+++ b/Dockerfile
@@ -5,9 +5,10 @@ ARG BASE_IMAGE
 FROM $BASE_IMAGE
 
 ARG DEP_GROUPS
+ARG BRANCH_NAME
 
 # Install and uninstall foundry to cache foundry requirements
-RUN git clone -b main https://github.com/mosaicml/llm-foundry.git
+RUN git clone -b $BRANCH_NAME https://github.com/mosaicml/llm-foundry.git
 RUN pip install --no-cache-dir --upgrade "./llm-foundry${DEP_GROUPS}"
 RUN pip uninstall -y llm-foundry
 RUN rm -rf llm-foundry