diff --git a/llama_stack/distribution/build_container.sh b/llama_stack/distribution/build_container.sh
index ba1863e5d..4924ad552 100755
--- a/llama_stack/distribution/build_container.sh
+++ b/llama_stack/distribution/build_container.sh
@@ -140,13 +140,17 @@ if command -v selinuxenabled &>/dev/null && selinuxenabled; then
   DOCKER_OPTS="$DOCKER_OPTS --security-opt label=disable"
 fi
 
+# Set version tag based on PyPI version
 if [ -n "$TEST_PYPI_VERSION" ]; then
-  image_name="$image_name-test-$TEST_PYPI_VERSION"
+  version_tag="test-$TEST_PYPI_VERSION"
 else
   URL="https://pypi.org/pypi/llama-stack/json"
-  image_name="$image_name-$(curl -s $URL | jq -r '.info.version')"
+  version_tag=$(curl -s $URL | jq -r '.info.version')
 fi
 
+# Add version tag to image name
+image_tag="$image_name:$version_tag"
+
 # Detect platform architecture
 ARCH=$(uname -m)
 if [ "$ARCH" = "arm64" ] || [ "$ARCH" = "aarch64" ]; then
diff --git a/llama_stack/distribution/start_container.sh b/llama_stack/distribution/start_container.sh
index b9ec9a23d..1efb76fb9 100755
--- a/llama_stack/distribution/start_container.sh
+++ b/llama_stack/distribution/start_container.sh
@@ -56,17 +56,18 @@ if [ -n "$LLAMA_CHECKPOINT_DIR" ]; then
   DOCKER_OPTS="$DOCKER_OPTS --gpus=all"
 fi
 
+version_tag="latest"
 if [ -n "$PYPI_VERSION" ]; then
-  docker_image="$docker_image-$PYPI_VERSION"
+  version_tag="$PYPI_VERSION"
 elif [ -n "$TEST_PYPI_VERSION" ]; then
-  docker_image="$docker_image-test-$TEST_PYPI_VERSION"
+  version_tag="test-$TEST_PYPI_VERSION"
 fi
 
 $DOCKER_BINARY run $DOCKER_OPTS -it \
   -p $port:$port \
   -v "$yaml_config:/app/config.yaml" \
   $mounts \
-  $docker_image \
+  $docker_image:$version_tag \
   python -m llama_stack.distribution.server.server \
   --yaml_config /app/config.yaml \
   --port $port "$@"