Torch upgrade

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/integ/sagemaker/experiments/helpers.py

@@ -44,15 +44,16 @@ def cleanup_exp_resources(exp_names, sagemaker_session):
             exp = Experiment.load(experiment_name=exp_name, sagemaker_session=sagemaker_session)
             exp._delete_all(action="--force")
 
+
 @pytest.fixture
 def clear_run_context():
     current_run = _RunContext.get_current_run()
-    if current_run == None:
+    if current_run is None:
         return
 
     logging.info(
         f"RunContext already populated by run {current_run.run_name}"
         f" in experiment {current_run.experiment_name}."
         " Clearing context manually"
     )
-    _RunContext.drop_current_run()
+    _RunContext.drop_current_run()

tests/integ/sagemaker/serve/test_serve_mlflow_pytorch_flavor_happy.py

@@ -166,9 +166,9 @@ def model_builder(request):
 #                 ), f"{caught_ex} was thrown when running pytorch squeezenet local container test"
 
 
-@pytest.mark.skipif(
-    PYTHON_VERSION_IS_NOT_310,  # or NOT_RUNNING_ON_INF_EXP_DEV_PIPELINE,
-    reason="The goal of these test are to test the serving components of our feature",
+@pytest.mark.skip(
+    reason="Testing against Python version 310 which is not supported anymore"
+    " https://github.com/aws/deep-learning-containers/blob/master/available_images.md",
 )
 def test_happy_pytorch_sagemaker_endpoint_with_torch_serve(
     sagemaker_session,

tests/unit/sagemaker/jumpstart/constants.py

@@ -17393,7 +17393,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",