docker debug

.github/workflows/docker.yaml

@@ -17,18 +17,9 @@ jobs:
     strategy:
       matrix:
         include:
-        - name: '2.1.0_cu121'
-          base_image: mosaicml/pytorch:2.1.0_cu121-python3.10-ubuntu20.04
-          dep_groups: '[gpu]'
         - name: '2.1.0_cu121_flash2'
           base_image: mosaicml/pytorch:2.1.0_cu121-python3.10-ubuntu20.04
           dep_groups: '[gpu-flash2]'
-        - name: '2.1.0_cu121_aws'
-          base_image: mosaicml/pytorch:2.1.0_cu121-python3.10-ubuntu20.04-aws
-          dep_groups: '[gpu]'
-        - name: '2.1.0_cu121_flash2_aws'
-          base_image: mosaicml/pytorch:2.1.0_cu121-python3.10-ubuntu20.04-aws
-          dep_groups: '[gpu-flash2]'
     steps:
     - name: Maximize Build Space on Worker
       uses: easimon/maximize-build-space@v4
@@ -81,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 }}