diff --git a/.github/workflows/release.yaml b/.github/workflows/release.yaml
index 6f1be030d7..47d5448c3d 100644
--- a/.github/workflows/release.yaml
+++ b/.github/workflows/release.yaml
@@ -10,46 +10,46 @@ jobs:
   code-quality:
     uses: ./.github/workflows/code-quality.yaml
 
-  # pypi-packaging:
-  #   name: Build and Publish llm-foundry PyPI Package
-  #   needs:
-  #   - code-quality
-  #   runs-on: linux-ubuntu-latest
-  #   steps:
-  #   - name: Checkout source
-  #     uses: actions/checkout@v3
-
-  #   - name: Set up Python
-  #     uses: actions/setup-python@v3
-  #     with:
-  #       python-version: "3.9"
-
-  #   - name: Build source and wheel distributions
-  #     run: |
-  #       if [[ "${{ github.ref }}" =~ refs\/tags\/v ]]; then
-  #         PYPI_PACKAGE_NAME="llm-foundry"
-  #       else
-  #         PYPI_PACKAGE_NAME="llm-foundry-test-$(date +%Y%m%d%H%M%S)"
-  #       fi
-
-  #       python -m pip install --upgrade build twine
-  #       python -m build
-  #       twine check --strict dist/*
-
-  #   - name: Publish 📦 to PyPI
-  #     uses: pypa/gh-action-pypi-publish@release/v1
-  #     if: contains(github.ref, 'refs/tags/v')
-  #     with:
-  #       user: __token__
-  #       password: ${{ secrets.PROD_PYPI_API_TOKEN }}
-
-  #   - name: Publish distribution 📦 to Test PyPI
-  #     uses: pypa/gh-action-pypi-publish@release/v1
-  #     if: contains(github.ref, 'refs/heads/') || contains(github.ref, 'refs/pull/')
-  #     with:
-  #       user: __token__
-  #       password: ${{ secrets.TEST_PYPI_API_TOKEN }}
-  #       repository_url: https://test.pypi.org/legacy/
+  pypi-packaging:
+    name: Build and Publish llm-foundry PyPI Package
+    needs:
+    - code-quality
+    runs-on: linux-ubuntu-latest
+    steps:
+    - name: Checkout source
+      uses: actions/checkout@v3
+
+    - name: Set up Python
+      uses: actions/setup-python@v3
+      with:
+        python-version: "3.9"
+
+    - name: Build source and wheel distributions
+      run: |
+        if [[ "${{ github.ref }}" =~ refs\/tags\/v ]]; then
+          PYPI_PACKAGE_NAME="llm-foundry"
+        else
+          PYPI_PACKAGE_NAME="llm-foundry-test-$(date +%Y%m%d%H%M%S)"
+        fi
+
+        python -m pip install --upgrade build twine
+        python -m build
+        twine check --strict dist/*
+
+    - name: Publish 📦 to PyPI
+      uses: pypa/gh-action-pypi-publish@release/v1
+      if: contains(github.ref, 'refs/tags/v')
+      with:
+        user: __token__
+        password: ${{ secrets.PROD_PYPI_API_TOKEN }}
+
+    - name: Publish distribution 📦 to Test PyPI
+      uses: pypa/gh-action-pypi-publish@release/v1
+      if: contains(github.ref, 'refs/heads/') || contains(github.ref, 'refs/pull/')
+      with:
+        user: __token__
+        password: ${{ secrets.TEST_PYPI_API_TOKEN }}
+        repository_url: https://test.pypi.org/legacy/
 
   build-docker:
     name: Build llm-foundry Release Docker Image
@@ -79,6 +79,8 @@ jobs:
 
         echo "DOCKER_TAG=mosaicml/llm-foundry:release_${TAG_NAME}" >> $GITHUB_ENV
         echo "AWS_DOCKER_TAG=mosaicml/llm-foundry:release_${TAG_NAME}_aws" >> $GITHUB_ENV
+        # echo "LATEST_TAG=mosaicml/llm-foundry:release-latest" >> $GITHUB_ENV
+        # echo "AWS_LATEST_TAG=mosaicml/llm-foundry:release_aws-latest" >> $GITHUB_ENV
 
 
     - name: Build and push AWS Docker image
@@ -89,6 +91,7 @@ jobs:
         push: true
         tags: |
           ${{ env.AWS_DOCKER_TAG }}
+          # ${{ env.AWS_LATEST_TAG }}
         build-args: |
           BASE_IMAGE=mosaicml/pytorch:2.4.0_cu124-python3.11-ubuntu20.04-aws
           BRANCH_NAME=${{ env.BRANCH_NAME }}
@@ -104,6 +107,7 @@ jobs:
         push: true
         tags: |
           ${{ env.DOCKER_TAG }}
+          # ${{ env.LATEST_TAG }}
         build-args: |
           BASE_IMAGE=mosaicml/pytorch:2.4.0_cu124-python3.11-ubuntu20.04
           BRANCH_NAME=${{ env.BRANCH_NAME }}
diff --git a/Dockerfile b/Dockerfile
index 5e8c43d3f3..80f2c82b10 100644
--- a/Dockerfile
+++ b/Dockerfile
@@ -9,6 +9,9 @@ ARG DEP_GROUPS
 ARG TE_COMMIT
 ARG KEEP_FOUNDRY=false
 
+RUN echo "KEEP_FOUNDRY is set to: '$KEEP_FOUNDRY'"
+ENV KEEP_FOUNDRY=${KEEP_FOUNDRY}
+
 ENV TORCH_CUDA_ARCH_LIST="8.0 8.6 8.7 8.9 9.0"
 
 # Check for changes in setup.py.