We read every piece of feedback, and take your input very seriously.
To see all available qualifiers, see our documentation.
1 parent 967c8d4 commit 331c7a1Copy full SHA for 331c7a1
src/sagemaker/image_uris.py
@@ -678,9 +678,9 @@ def get_training_image_uri(
678
if "modelparallel" in distribution["smdistributed"]:
679
if distribution["smdistributed"]["modelparallel"].get("enabled", True):
680
framework = "pytorch-smp"
681
- if "p5" in instance_type
682
- or "2.1" in framework_version
683
- or "2.2" in framework_version:
+ if "p5" in instance_type or \
+ "2.1" in framework_version or \
+ "2.2" in framework_version:
684
container_version = "cu121"
685
else:
686
container_version = "cu118"
0 commit comments