Skip to content

Commit c4fb2a2

Browse files
author
wanyixia
committed
change: add compilation option for ml_eia2
1 parent 999fb22 commit c4fb2a2

File tree

1 file changed

+3
-1
lines changed

1 file changed

+3
-1
lines changed

src/sagemaker/model.py

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -651,7 +651,9 @@ def compile(
651651
job_status = self.sagemaker_session.wait_for_compilation_job(job_name)
652652
self.model_data = job_status["ModelArtifacts"]["S3ModelArtifacts"]
653653
if target_instance_family is not None:
654-
if target_instance_family.startswith("ml_"):
654+
if target_instance_family == "ml_eia2":
655+
LOGGER.info("You are using target device ml_eia2...")
656+
elif (target_instance_family != "ml_eia2") and target_instance_family.startswith("ml_"):
655657
self.image_uri = self._compilation_image_uri(
656658
self.sagemaker_session.boto_region_name,
657659
target_instance_family,

0 commit comments

Comments
 (0)