Revert "Fix torch 2.10 version pinning and increase CPU image size baseline"

This reverts commit 73dfc44.

Author: Bhanu Teja Goshikonda

pytorch/training/buildspec-2-10-ec2.yml

@@ -44,7 +44,7 @@ images:
   BuildEC2CPUPTTrainPy3DockerImage:
     <<: *TRAINING_REPOSITORY
     build: &PYTORCH_CPU_TRAINING_PY3 false
-    image_size_baseline: 12000
+    image_size_baseline: 7200
     device_type: &DEVICE_TYPE cpu
     python_version: &DOCKER_PYTHON_VERSION py3
     tag_python_version: &TAG_PYTHON_VERSION py313

pytorch/training/buildspec-2-10-sm.yml

@@ -44,7 +44,7 @@ images:
   BuildSageMakerCPUPTTrainPy3DockerImage:
     <<: *TRAINING_REPOSITORY
     build: &PYTORCH_CPU_TRAINING_PY3 false
-    image_size_baseline: 12000
+    image_size_baseline: 7200
     device_type: &DEVICE_TYPE cpu
     python_version: &DOCKER_PYTHON_VERSION py3
     tag_python_version: &TAG_PYTHON_VERSION py313

pytorch/training/docker/2.10/py3/Dockerfile.cpu

@@ -197,26 +197,20 @@ RUN pip install --no-cache-dir \
     pytz \
     tzdata
 
-# Install PyTorch - split into separate commands to prevent dependency resolver from downgrading torch
+# Install PyTorch
 RUN pip install --no-cache-dir -U torch==${PYTORCH_VERSION} \
     torchvision==${TORCHVISION_VERSION} \
     torchaudio==${TORCHAUDIO_VERSION} \
-    --index-url https://download.pytorch.org/whl/cpu
-
-# Install torch ecosystem packages
-RUN pip install --no-cache-dir -U torchtnt==${TORCHTNT_VERSION} \
+    --index-url https://download.pytorch.org/whl/cpu \
+    && pip install --no-cache-dir -U torchtnt==${TORCHTNT_VERSION} \
     torchdata==${TORCHDATA_VERSION} \
-    s3torchconnector
-
-# Install ML packages with torch version constraint to prevent downgrade
-# pin numpy requirement for fastai dependency
-# requires explicit declaration of spacy, thinc, blis
-# pin thinc due to incompatibility with numpy 1.26.4 (sagemaker doesn't support latest numpy)
-RUN pip install --no-cache-dir -U \
-    "torch==${PYTORCH_VERSION}" \
+    s3torchconnector \
     fastai \
     accelerate \
+    # pin numpy requirement for fastai dependency
+    # requires explicit declaration of spacy, thic, blis
     spacy \
+    # pin thinc due to incompatibility with numpy 1.26.4 (sagemaker doesn't support latest numpy)
     thinc \
     blis \
     numpy \

pytorch/training/docker/2.10/py3/cu130/Dockerfile.gpu

@@ -116,27 +116,21 @@ RUN pip install --no-cache-dir \
     pytz \
     tzdata
 
-# Install PyTorch - split into separate commands to prevent dependency resolver from downgrading torch
+# Install PyTorch
 RUN pip install --no-cache-dir -U torch==${PYTORCH_VERSION} \
     torchvision==${TORCHVISION_VERSION} \
     torchaudio==${TORCHAUDIO_VERSION} \
-    --index-url https://download.pytorch.org/whl/cu130
-
-# Install torch ecosystem packages
-RUN pip install --no-cache-dir -U torchtnt==${TORCHTNT_VERSION} \
+    --index-url https://download.pytorch.org/whl/cu130 \
+    && pip install --no-cache-dir -U torchtnt==${TORCHTNT_VERSION} \
     torchdata==${TORCHDATA_VERSION} \
     triton \
-    s3torchconnector
-
-# Install ML packages with torch version constraint to prevent downgrade
-# pin numpy requirement for fastai dependency
-# requires explicit declaration of spacy, thinc, blis
-# pin thinc due to incompatibility with numpy 1.26.4 (sagemaker doesn't support latest numpy)
-RUN pip install --no-cache-dir -U \
-    "torch==${PYTORCH_VERSION}" \
+    s3torchconnector \
     fastai \
     accelerate \
+    # pin numpy requirement for fastai dependency
+    # requires explicit declaration of spacy, thic, blis
     spacy \
+    # pin thinc due to incompatibility with numpy 1.26.4 (sagemaker doesn't support latest numpy)
     thinc \
     blis \
     numpy \