UPDATE PYTORCH VERSION TO ADDRESS SECURITY RISK

**Description**

Currently used Pytorch version has a possible vulnerability .

Internal - https://tiny.amazon.com/p5i4jla1

**Testing Done**

Unit and Integration tests in the CodeBuild

Author: nargokul

tests/data/modules/script_mode/requirements.txt

@@ -1,3 +1,3 @@
 numpy
 -f https://download.pytorch.org/whl/torch_stable.html
-torch==2.0.1+cpu
+torch>=2.6.0+cpu

tests/data/serve_resources/mlflow/pytorch/conda.yaml

@@ -17,8 +17,8 @@ dependencies:
   - pandas==2.2.1
   - pyyaml==6.0.1
   - requests==2.31.0
-  - torch==2.0.1
-  - torchvision==0.15.2
+  - torch>=2.6.0
+  - torchvision>=0.17.0
   - tqdm==4.66.2
   - scikit-learn==1.3.2
 name: mlflow-env

tests/data/serve_resources/mlflow/pytorch/requirements.txt

@@ -11,6 +11,6 @@ packaging==21.3
 pandas==2.2.1
 pyyaml==6.0.1
 requests==2.32.2
-torch==2.2.0
-torchvision==0.17.0
+torch>=2.6.0
+torchvision>=0.17.0
 tqdm==4.66.3

tests/unit/sagemaker/jumpstart/constants.py

@@ -17391,7 +17391,7 @@
             "texttable==1.6.7",
             "tokenize-rt==5.1.0",
             "tokenizers==0.13.3",
-            "torch==2.2.0",
+            "torch>=2.6.0",
             "transformers==4.33.3",
             "triton==2.2.0",
             "typing-extensions==4.8.0",

tox.ini

@@ -83,8 +83,8 @@ passenv =
 commands =
     python -c "import os; os.system('install-custom-pkgs --install-boto-wheels')"
     pip install 'apache-airflow==2.9.3' --constraint "https://raw.githubusercontent.com/apache/airflow/constraints-2.9.3/constraints-3.8.txt"
-    pip install 'torch==2.0.1+cpu' -f 'https://download.pytorch.org/whl/torch_stable.html'
-    pip install 'torchvision==0.15.2+cpu' -f 'https://download.pytorch.org/whl/torch_stable.html'
+    pip install 'torch>=2.6.0+cpu' -f 'https://download.pytorch.org/whl/torch_stable.html'
+    pip install 'torchvision>=0.17.0+cpu' -f 'https://download.pytorch.org/whl/torch_stable.html'
     pip install 'dill>=0.3.8'
 
     pytest {posargs}