fix tests

aws · Oct 17, 2024 · 9374c75 · 9374c75
1 parent 815f1a4
commit 9374c75
Showing 1 changed file with 1 addition and 1 deletion.
diff --git 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",