Skip to content

Commit 6211227

Browse files
committed
Address PR comments
1 parent 43ce1ba commit 6211227

File tree

1 file changed

+2
-4
lines changed

1 file changed

+2
-4
lines changed

src/sagemaker/serve/builder/transformers_builder.py

Lines changed: 2 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -100,11 +100,9 @@ def _create_transformers_model(self) -> Type[Model]:
100100
if hf_model_md is None:
101101
raise ValueError("Could not fetch HF metadata")
102102

103-
model_task = None
104-
if self.model_metadata:
103+
model_task = hf_model_md.get("pipeline_tag")
104+
if self.model_metadata.get("HF_TASK") is not None:
105105
model_task = self.model_metadata.get("HF_TASK")
106-
else:
107-
model_task = hf_model_md.get("pipeline_tag")
108106

109107
if model_task == "sentence-similarity" and not self.image_uri:
110108
self.image_uri = get_huggingface_llm_image_uri(

0 commit comments

Comments
 (0)