diff --git a/src/sagemaker/image_uri_config/huggingface-llm.json b/src/sagemaker/image_uri_config/huggingface-llm.json index ca5b05d3eb..b3988362fd 100644 --- a/src/sagemaker/image_uri_config/huggingface-llm.json +++ b/src/sagemaker/image_uri_config/huggingface-llm.json @@ -717,7 +717,7 @@ "tag_prefix": "2.3.0-tgi2.2.0", "repository": "huggingface-pytorch-tgi-inference", "container_version": { - "gpu": "cu121-ubuntu22.04" + "gpu": "cu121-ubuntu22.04-v2.0" } } } diff --git a/tests/unit/sagemaker/image_uris/test_huggingface_llm.py b/tests/unit/sagemaker/image_uris/test_huggingface_llm.py index e414455d17..e753687018 100644 --- a/tests/unit/sagemaker/image_uris/test_huggingface_llm.py +++ b/tests/unit/sagemaker/image_uris/test_huggingface_llm.py @@ -42,7 +42,7 @@ "2.0.0": "2.1.1-tgi2.0.0-gpu-py310-cu121-ubuntu22.04", "2.0.1": "2.1.1-tgi2.0.1-gpu-py310-cu121-ubuntu22.04", "2.0.2": "2.3.0-tgi2.0.2-gpu-py310-cu121-ubuntu22.04", - "2.2.0": "2.3.0-tgi2.2.0-gpu-py310-cu121-ubuntu22.04", + "2.2.0": "2.3.0-tgi2.2.0-gpu-py310-cu121-ubuntu22.04-v2.0", }, "inf2": { "0.0.16": "1.13.1-optimum0.0.16-neuronx-py310-ubuntu22.04",