diff --git a/02-ollama_install.sh b/02-ollama_install.sh
index 13ec0a4..bbdb526 100644
--- a/02-ollama_install.sh
+++ b/02-ollama_install.sh
@@ -3,7 +3,11 @@ path=$PWD
 
 sudo apt update
 sudo apt upgrade -y
+sudo apt install -y linux-headers-$(uname -r) build-essential dkms
+
 sudo ubuntu-drivers install --gpgpu
+sudo apt install -y nvidia-container-toolkit
+sudo nvidia-ctk runtime configure --runtime=docker && sudo systemctl restart docker
 # install ollama
 sudo curl -fsSL https://ollama.com/install.sh | sh