Revert "bump tests no transform"

This reverts commit bc999ad.
mosaicml · Jan 13, 2024 · 6b0b0ef · 6b0b0ef
1 parent bc999ad
commit 6b0b0ef
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/docker/Dockerfile b/docker/Dockerfile
@@ -305,6 +305,7 @@ RUN if [ -n "$CUDA_VERSION" ] ; then \
         cd .. && \
         pip install --no-dependencies "git+ssh://[email protected]/mosaicml/megablocks-private.git@main#egg=megablocks[all]" && \
         pip install --no-dependencies git+https://github.com/tgale96/[email protected] && \
+        pip install git+https://github.com/NVIDIA/TransformerEngine.git@stable && \
         pip uninstall -y torch && \
         pip uninstall -y pytorch-triton && \
         pip uninstall -y torchvision && \