fix: p2 issue for test_pytorch_training

Author: Shibo Xing

tests/unit/sagemaker/image_uris/test_dlc_frameworks.py

@@ -18,6 +18,7 @@
 from tests.unit.sagemaker.image_uris import expected_uris
 
 INSTANCE_TYPES_AND_PROCESSORS = (("ml.c4.xlarge", "cpu"), ("ml.p2.xlarge", "gpu"))
+RENEWED_PYTORCH_INSTANCE_TYPES_AND_PROCESSORS = (("ml.c4.xlarge", "cpu"), ("ml.g4dn.xlarge", "gpu"))
 REGION = "us-west-2"
 
 DLC_ACCOUNT = "763104351884"
@@ -70,7 +71,12 @@ def _test_image_uris(
         "image_scope": scope,
     }
 
-    for instance_type, processor in INSTANCE_TYPES_AND_PROCESSORS:
+    TYPES_AND_PROCESSORS = INSTANCE_TYPES_AND_PROCESSORS
+    if framework == "pytorch" and Version(fw_version) >= Version("1.13"):
+        '''Handle P2 deprecation'''
+        TYPES_AND_PROCESSORS = RENEWED_PYTORCH_INSTANCE_TYPES_AND_PROCESSORS
+
+    for instance_type, processor in TYPES_AND_PROCESSORS:
         uri = image_uris.retrieve(region=REGION, instance_type=instance_type, **base_args)
 
         expected = expected_fn(processor=processor, **expected_fn_args)