fix branch name

.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 }}

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