Update cu121

Author: Teng-xu

src/sagemaker/image_uris.py

@@ -667,7 +667,7 @@ def get_training_image_uri(
                 if distribution["smdistributed"]["modelparallel"].get("enabled", True):
                     framework = "pytorch-smp"
                 if "p5" in instance_type:
-                    container_version = "cu12"
+                    container_version = "cu121"
                 else:
                     container_version = "cu118"
 

tests/unit/sagemaker/image_uris/test_smp_v2.py

@@ -16,7 +16,7 @@
 from sagemaker import image_uris
 from tests.unit.sagemaker.image_uris import expected_uris
 
-CONTAINER_VERSIONS = {"ml.p4d.24xlarge": "cu118", "ml.p5d.24xlarge": "cu12"}
+CONTAINER_VERSIONS = {"ml.p4d.24xlarge": "cu118", "ml.p5d.24xlarge": "cu121"}
 
 
 @pytest.mark.parametrize("load_config", ["pytorch-smp.json"], indirect=True)