Corrections

Author: saimidu

src/sagemaker/image_uri_config/pytorch.json

@@ -890,7 +890,7 @@
                     "us-west-1": "763104351884",
                     "us-west-2": "763104351884"
                 },
-                "repository": "pytorch-inference"
+                "repository": "pytorch-training"
             }
         }
     }

tests/conftest.py

@@ -186,6 +186,8 @@ def pytorch_training_py_version(pytorch_training_version, request):
         return request.param
     elif Version(pytorch_training_version) == Version("1.7.1"):
         return "py36"
+    elif Version(pytorch_training_version) >= Version("1.9"):
+        return "py38"
     else:
         return "py3"
 
@@ -196,6 +198,8 @@ def pytorch_inference_py_version(pytorch_inference_version, request):
         return request.param
     elif Version(pytorch_inference_version) == Version("1.7.1"):
         return "py36"
+    elif Version(pytorch_inference_version) >= Version("1.9"):
+        return "py38"
     else:
         return "py3"
 

tests/unit/test_processing.py

@@ -284,14 +284,20 @@ def test_sklearn_with_all_parameters_via_run_args_called_twice(
 @patch("os.path.exists", return_value=True)
 @patch("os.path.isfile", return_value=True)
 def test_pytorch_processor_with_required_parameters(
-    exists_mock, isfile_mock, botocore_resolver, sagemaker_session, pytorch_training_version
+    exists_mock,
+    isfile_mock,
+    botocore_resolver,
+    sagemaker_session,
+    pytorch_training_version,
+    pytorch_training_py_version,
 ):
     botocore_resolver.return_value.construct_endpoint.return_value = {"hostname": ECR_HOSTNAME}
 
     processor = PyTorchProcessor(
         role=ROLE,
         instance_type="ml.m4.xlarge",
         framework_version=pytorch_training_version,
+        py_version=pytorch_training_py_version,
         instance_count=1,
         sagemaker_session=sagemaker_session,
     )