From 9374c7518e97bd845f952f923d8784cbedb02f02 Mon Sep 17 00:00:00 2001
From: Haotian An <haotiaa@amazon.com>
Date: Thu, 17 Oct 2024 01:26:52 +0000
Subject: [PATCH] fix tests

---
 tests/unit/sagemaker/image_uris/test_huggingface_llm.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tests/unit/sagemaker/image_uris/test_huggingface_llm.py b/tests/unit/sagemaker/image_uris/test_huggingface_llm.py
index 581dcf90c1..28525a390c 100644
--- a/tests/unit/sagemaker/image_uris/test_huggingface_llm.py
+++ b/tests/unit/sagemaker/image_uris/test_huggingface_llm.py
@@ -45,7 +45,7 @@
         "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-v2.0",
-        "2.3.1": "2.4.0-tgi2.3.1-gpu-py310-cu124-ubuntu22.04",
+        "2.3.1": "2.4.0-tgi2.3.1-gpu-py311-cu124-ubuntu22.04",
     },
     "inf2": {
         "0.0.16": "1.13.1-optimum0.0.16-neuronx-py310-ubuntu22.04",