Address PR comments

aws · May 15, 2024 · 6211227 · 6211227
1 parent 43ce1ba
commit 6211227
Showing 1 changed file with 2 additions and 4 deletions.
diff --git a/src/sagemaker/serve/builder/transformers_builder.py b/src/sagemaker/serve/builder/transformers_builder.py
@@ -100,11 +100,9 @@ def _create_transformers_model(self) -> Type[Model]:
         if hf_model_md is None:
             raise ValueError("Could not fetch HF metadata")
 
-        model_task = None
-        if self.model_metadata:
+        model_task = hf_model_md.get("pipeline_tag")
+        if self.model_metadata.get("HF_TASK") is not None:
             model_task = self.model_metadata.get("HF_TASK")
-        else:
-            model_task = hf_model_md.get("pipeline_tag")
 
         if model_task == "sentence-similarity" and not self.image_uri:
             self.image_uri = get_huggingface_llm_image_uri(