From 4254378c0fca96921299a1f0d736509a2e58c415 Mon Sep 17 00:00:00 2001 From: wayneyao Date: Thu, 30 Nov 2023 15:57:15 -0800 Subject: [PATCH] bug fix: save utils fix --- src/sagemaker/serve/save_retrive/version_1_0_0/save/utils.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/sagemaker/serve/save_retrive/version_1_0_0/save/utils.py b/src/sagemaker/serve/save_retrive/version_1_0_0/save/utils.py index 10a05a5307..7370066415 100644 --- a/src/sagemaker/serve/save_retrive/version_1_0_0/save/utils.py +++ b/src/sagemaker/serve/save_retrive/version_1_0_0/save/utils.py @@ -186,5 +186,5 @@ def get_image_url(model_server: str) -> str: "DJLServing": "763104351884.dkr.ecr.us-west-2.amazonaws.com/djl-inference:0.25.0-cpu-full" } if model_server in supported_image: - return supported_image.get("model_server") + return supported_image.get(model_server) raise Exception("Model server specified is not supported in SageMaker save format")