fix: save utils fix (#4293)

haNa-meister · wayneyao · web-flow · commit 57fd632a1b4b · 2023-12-01T18:25:27.000-08:00
Co-authored-by: wayneyao &lt;wayneyao@amazon.com&gt;
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
@@ -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")

Original file line number	Diff line number	Diff line change
`@@ -186,5 +186,5 @@ def get_image_url(model_server: str) -> str:`
`186`	`186`	`"DJLServing": "763104351884.dkr.ecr.us-west-2.amazonaws.com/djl-inference:0.25.0-cpu-full"`
`187`	`187`	`}`
`188`	`188`	`if model_server in supported_image:`
`189`		`- return supported_image.get("model_server")`
	`189`	`+ return supported_image.get(model_server)`
`190`	`190`	`raise Exception("Model server specified is not supported in SageMaker save format")`