changing logic around for unit test

Author: adtian2

tests/unit/sagemaker/image_uris/test_smp_v2.py

@@ -32,15 +32,16 @@ def test_smp_v2(load_config):
             ACCOUNTS = load_config["training"]["versions"][version]["registries"]
             PY_VERSIONS = load_config["training"]["versions"][version]["py_versions"]
             for py_version in PY_VERSIONS:
+                # py311 is only for smp version 2.4.0
+                if py_version == "py311" and "2.4" not in version:
+                    continue
+
                 for region in ACCOUNTS.keys():
                     for instance_type in CONTAINER_VERSIONS.keys():
                         cuda_vers = CONTAINER_VERSIONS[instance_type]
                         if "2.1" in version or "2.2" in version or "2.3" in version:
                             cuda_vers = "cu121"
 
-                        if "2.4" in version:
-                            py_version = "py311"
-
                         uri = image_uris.get_training_image_uri(
                             region,
                             framework="pytorch",