diff --git a/docker/Dockerfile b/docker/Dockerfile
index cc55ed0c84..6c9dcd5da6 100644
--- a/docker/Dockerfile
+++ b/docker/Dockerfile
@@ -49,9 +49,6 @@ ARG URLLIB3_VERSION='>=1.26.5,<2'
 # Build the composer image on the pytorch image
 ARG COMPOSER_BASE=pytorch_stage
 
-# The command that is passed to `pip install` -- e.g. `pip install "${COMPOSER_INSTALL_COMMAND}"`
-ARG COMPOSER_INSTALL_COMMAND='mosaicml[all]'
-
 #########################
 # Build the PyTorch Image
 #########################
@@ -343,4 +340,7 @@ ARG DEBIAN_FRONTEND=noninteractive
 # Install Composer
 ##################
 
+# The command that is passed to `pip install` -- e.g. `pip install "${COMPOSER_INSTALL_COMMAND}"`
+ARG COMPOSER_INSTALL_COMMAND='mosaicml[all]'
+
 RUN pip install "${COMPOSER_INSTALL_COMMAND}"