test

mosaicml · Oct 19, 2024 · 9728a7c · 9728a7c
1 parent 3a6206a
commit 9728a7c
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/Dockerfile b/Dockerfile
@@ -2,6 +2,7 @@ ARG BASE_IMAGE
 FROM $BASE_IMAGE
 
 # Install TransformerEngine
+ARG TE_COMMIT
 RUN NVTE_FRAMEWORK=pytorch CMAKE_BUILD_PARALLEL_LEVEL=4 MAX_JOBS=4 pip install git+https://github.com/NVIDIA/TransformerEngine.git@$TE_COMMIT
 
 ARG BRANCH_NAME