fixing formatting

Author: Andrew Tian

src/sagemaker/fw_utils.py

@@ -161,7 +161,14 @@
 ]
 
 
-TORCH_DISTRIBUTED_GPU_SUPPORTED_FRAMEWORK_VERSIONS = ["1.13.1", "2.0.0", "2.0.1", "2.1.0", "2.1.2", "2.2.0"]
+TORCH_DISTRIBUTED_GPU_SUPPORTED_FRAMEWORK_VERSIONS = [
+    "1.13.1",
+    "2.0.0",
+    "2.0.1",
+    "2.1.0",
+    "2.1.2",
+    "2.2.0",
+]
 
 TRAINIUM_SUPPORTED_DISTRIBUTION_STRATEGIES = ["torch_distributed"]
 TRAINIUM_SUPPORTED_TORCH_DISTRIBUTED_FRAMEWORK_VERSIONS = [

src/sagemaker/image_uris.py

@@ -678,9 +678,11 @@ def get_training_image_uri(
             if "modelparallel" in distribution["smdistributed"]:
                 if distribution["smdistributed"]["modelparallel"].get("enabled", True):
                     framework = "pytorch-smp"
-                    if "p5" in instance_type or \
-                        "2.1" in framework_version or \
-                        "2.2" in framework_version:
+                    if (
+                        "p5" in instance_type
+                        or "2.1" in framework_version
+                        or "2.2" in framework_version
+                    ):
                         container_version = "cu121"
                     else:
                         container_version = "cu118"