From 0861539a49bb14589aa5bba65d6917879e2951f0 Mon Sep 17 00:00:00 2001
From: kshitizgupta21 <73861554+kshitizgupta21@users.noreply.github.com>
Date: Tue, 19 Nov 2024 10:47:07 -0800
Subject: [PATCH] update inference ami version in sagemaker endpoint config to
 fix nvml driver issues (#104)

---
 cloud-service-providers/aws/sagemaker/nim_llama3.ipynb | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/cloud-service-providers/aws/sagemaker/nim_llama3.ipynb b/cloud-service-providers/aws/sagemaker/nim_llama3.ipynb
index 1b76d46..e8a64f1 100644
--- a/cloud-service-providers/aws/sagemaker/nim_llama3.ipynb
+++ b/cloud-service-providers/aws/sagemaker/nim_llama3.ipynb
@@ -329,7 +329,8 @@
     "            \"InitialInstanceCount\": 1,\n",
     "            \"ModelName\": sm_model_name,\n",
     "            \"VariantName\": \"AllTraffic\",\n",
-    "            \"ContainerStartupHealthCheckTimeoutInSeconds\": 850\n",
+    "            \"ContainerStartupHealthCheckTimeoutInSeconds\": 1800,\n",
+    "            \"InferenceAmiVersion\": \"al2-ami-sagemaker-inference-gpu-2\"\n",
     "        }\n",
     "    ],\n",
     ")\n",