diff --git a/.github/workflows/docker.yaml b/.github/workflows/docker.yaml
index 1b02b4fd9b..6ca10fcd47 100644
--- a/.github/workflows/docker.yaml
+++ b/.github/workflows/docker.yaml
@@ -12,7 +12,7 @@ on:
   workflow_dispatch: {}
 jobs:
   docker-build:
-    runs-on: mosaic-4wide
+    runs-on: mosaic-8wide
     if: github.repository_owner == 'mosaicml'
     strategy:
       matrix:
diff --git a/Dockerfile b/Dockerfile
index fd5b8e40ab..db68fc9d2b 100644
--- a/Dockerfile
+++ b/Dockerfile
@@ -13,7 +13,8 @@ ADD https://raw.githubusercontent.com/mosaicml/llm-foundry/$BRANCH_NAME/setup.py
 RUN rm setup.py
 
 # Install and uninstall foundry to cache foundry requirements
-RUN git clone -b $BRANCH_NAME https://github.com/mosaicml/llm-foundry.git
-RUN pip install --no-cache-dir "./llm-foundry${DEP_GROUPS}"
+RUN pip install --verbose --no-cache-dir "./llm-foundry${DEP_GROUPS}"
+RUN NVTE_FRAMEWORK=pytorch pip install --verbose --no-cache-dir git+https://github.com/NVIDIA/TransformerEngine.git@main
 RUN pip uninstall -y llm-foundry
+RUN pip uninstall -y transformer-engine
 RUN rm -rf llm-foundry