Skip to content

Commit 5995b27

Browse files
HappyAmazonianmufiAmazon
authored andcommitted
fix: fix ml_inf to ml_inf1 in Neo multi-version support (aws#3405)
1 parent 6bcb5a0 commit 5995b27

File tree

2 files changed

+4
-4
lines changed

2 files changed

+4
-4
lines changed

src/sagemaker/model.py

+1-1
Original file line numberDiff line numberDiff line change
@@ -789,7 +789,7 @@ def multi_version_compilation_supported(
789789
}
790790
if target_instance_type in NEO_IOC_TARGET_DEVICES:
791791
return framework in multi_version_frameworks_support_mapping["neo_ioc_targets"]
792-
if target_instance_type == "ml_inf":
792+
if target_instance_type == "ml_inf1":
793793
return framework in multi_version_frameworks_support_mapping["inferentia"]
794794
return False
795795

tests/unit/sagemaker/model/test_neo.py

+3-3
Original file line numberDiff line numberDiff line change
@@ -47,7 +47,7 @@ def test_compile_model_for_inferentia(sagemaker_session):
4747
)
4848
model = _create_model(sagemaker_session)
4949
model.compile(
50-
target_instance_family="ml_inf",
50+
target_instance_family="ml_inf1",
5151
input_shape={"data": [1, 3, 1024, 1024]},
5252
output_path="s3://output",
5353
role="role",
@@ -313,7 +313,7 @@ def test_compile_with_pytorch_neo_in_ml_inf(session):
313313

314314
model = _create_model()
315315
model.compile(
316-
target_instance_family="ml_inf",
316+
target_instance_family="ml_inf1",
317317
input_shape={"data": [1, 3, 1024, 1024]},
318318
output_path="s3://output",
319319
role="role",
@@ -336,7 +336,7 @@ def test_compile_with_tensorflow_neo_in_ml_inf(session):
336336

337337
model = _create_model()
338338
model.compile(
339-
target_instance_family="ml_inf",
339+
target_instance_family="ml_inf1",
340340
input_shape={"data": [1, 3, 1024, 1024]},
341341
output_path="s3://output",
342342
role="role",

0 commit comments

Comments
 (0)