fix unit tests

haixiw · haixiw · commit 815ae086cdf7 · 2024-05-20T17:52:17.000Z
diff --git a/tests/unit/sagemaker/image_uris/test_huggingface_llm.py b/tests/unit/sagemaker/image_uris/test_huggingface_llm.py
@@ -22,6 +22,9 @@
     "gpu": {
         "1.2.3": "2.0.1-tei1.2.3-gpu-py310-cu122-ubuntu22.04",
     },
+    "cpu": {
+        "1.2.3": "2.0.1-tei1.2.3-cpu-py310-ubuntu22.04",
+    },
 }
 HF_VERSIONS_MAPPING = {
     "gpu": {
