We read every piece of feedback, and take your input very seriously.
To see all available qualifiers, see our documentation.
1 parent 6bcb5a0 commit 5995b27Copy full SHA for 5995b27
src/sagemaker/model.py
@@ -789,7 +789,7 @@ def multi_version_compilation_supported(
789
}
790
if target_instance_type in NEO_IOC_TARGET_DEVICES:
791
return framework in multi_version_frameworks_support_mapping["neo_ioc_targets"]
792
- if target_instance_type == "ml_inf":
+ if target_instance_type == "ml_inf1":
793
return framework in multi_version_frameworks_support_mapping["inferentia"]
794
return False
795
tests/unit/sagemaker/model/test_neo.py
@@ -47,7 +47,7 @@ def test_compile_model_for_inferentia(sagemaker_session):
47
)
48
model = _create_model(sagemaker_session)
49
model.compile(
50
- target_instance_family="ml_inf",
+ target_instance_family="ml_inf1",
51
input_shape={"data": [1, 3, 1024, 1024]},
52
output_path="s3://output",
53
role="role",
@@ -313,7 +313,7 @@ def test_compile_with_pytorch_neo_in_ml_inf(session):
313
314
model = _create_model()
315
316
317
318
319
@@ -336,7 +336,7 @@ def test_compile_with_tensorflow_neo_in_ml_inf(session):
336
337
338
339
340
341
342
0 commit comments