diff --git a/LICENSE b/LICENSE
index 6775bd7..430a663 100644
--- a/LICENSE
+++ b/LICENSE
@@ -1,6 +1,7 @@
 MIT License
 
 Copyright (c) 2022 Microsoft Corporation
+Copyright (c) 2024 Alex Fraser
 
 Permission is hereby granted, free of charge, to any person obtaining a copy
 of this software and associated documentation files (the "Software"), to deal
diff --git a/src/ollama/devcontainer-feature.json b/src/ollama/devcontainer-feature.json
index ae12e53..e26d9c8 100644
--- a/src/ollama/devcontainer-feature.json
+++ b/src/ollama/devcontainer-feature.json
@@ -1,16 +1,18 @@
 {
     "name": "ollama",
     "id": "ollama",
-    "version": "1.0.0",
+    "version": "1.0.1",
     "description": "Installs ollama",
     "documentationURL": "https://github.com/prulloac/devcontainer-features/tree/main/src/ollama",
     "options": {
         "pull": {
-            "type":"string",
+            "type": "string",
             "default": "",
-            "description": "comma separated list of models to pull"
+            "description": "Comma-separated list of models to pull",
+            "proposals": ["llama3.2", "llama3.2:3b,nomic-embed-text"]
         }
     },
+    "onCreateCommand": "/usr/local/share/ollama/init.sh",
     "entrypoint": "ollama serve",
     "installsAfter": [
         "ghcr.io/devcontainers/features/common-utils"
diff --git a/src/ollama/install.sh b/src/ollama/install.sh
index a63375c..d26d06d 100644
--- a/src/ollama/install.sh
+++ b/src/ollama/install.sh
@@ -20,12 +20,19 @@ check_packages ca-certificates curl
 
 curl https://ollama.ai/install.sh | sh
 
-# if PULL is not empty and not equal to "none", also ollama is serving, then pull the models
-if [ "$PULL" != "none" ] && [ "$PULL" != "" ]; then
-    # split PULL variable into array using comma as delimiter and then pass to ollama pull
-    ollama serve &
-    sleep 3
-    echo $PULL | tr ',' '\n' | xargs -I % sh -c "ollama pull %"
-fi
-
-echo 'Done!'
\ No newline at end of file
+# Make sure it can run
+ollama --version
+
+# Save requested models for later. They can't be installed now because the user may be incorrect.
+mkdir -p "/usr/local/share/ollama"
+echo "$PULL" | tr ',' '\n' > "/usr/local/share/ollama/models.txt"
+cat << "EOF" > "/usr/local/share/ollama/init.sh"
+#!/bin/sh
+set -e
+echo "Pulling models" 1>&2
+# Ollama should already be running
+cat /usr/local/share/ollama/models.txt | xargs -I % sh -c "ollama pull %"
+EOF
+chmod +x "/usr/local/share/ollama/init.sh"
+
+echo 'Done!'