Skip to content

Commit 22cfd1a

Browse files
authored
feat: Add inf2 support to HuggingFaceModel (#3900)
1 parent 8794f8a commit 22cfd1a

File tree

1 file changed

+2
-0
lines changed

1 file changed

+2
-0
lines changed

src/sagemaker/huggingface/model.py

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -303,9 +303,11 @@ def deploy(
303303
"""
304304

305305
if not self.image_uri and instance_type is not None and instance_type.startswith("ml.inf"):
306+
inference_tool = "neuron" if instance_type.startswith("ml.inf1") else "neuronx"
306307
self.image_uri = self.serving_image_uri(
307308
region_name=self.sagemaker_session.boto_session.region_name,
308309
instance_type=instance_type,
310+
inference_tool=inference_tool,
309311
)
310312

311313
return super(HuggingFaceModel, self).deploy(

0 commit comments

Comments
 (0)