From 6b0b0efe0ec7500c52eb4ab2a3a9c2c2fa344087 Mon Sep 17 00:00:00 2001 From: Chuck Tang Date: Sat, 13 Jan 2024 00:20:35 -0800 Subject: [PATCH] Revert "bump tests no transform" This reverts commit bc999ad0aed07eb3d4e41aa88788f94d4d66138e. --- docker/Dockerfile | 1 + 1 file changed, 1 insertion(+) diff --git a/docker/Dockerfile b/docker/Dockerfile index e66c3a9170..9b604ad4fc 100644 --- a/docker/Dockerfile +++ b/docker/Dockerfile @@ -305,6 +305,7 @@ RUN if [ -n "$CUDA_VERSION" ] ; then \ cd .. && \ pip install --no-dependencies "git+ssh://git@github.com/mosaicml/megablocks-private.git@main#egg=megablocks[all]" && \ pip install --no-dependencies git+https://github.com/tgale96/grouped_gemm@v0.1.4 && \ + pip install git+https://github.com/NVIDIA/TransformerEngine.git@stable && \ pip uninstall -y torch && \ pip uninstall -y pytorch-triton && \ pip uninstall -y torchvision && \