diff --git a/tools/cli/commands/create.py b/tools/cli/commands/create.py
index dad07ed41..db763b409 100644
--- a/tools/cli/commands/create.py
+++ b/tools/cli/commands/create.py
@@ -296,6 +296,8 @@
        --env='DATALAB_GIT_AUTHOR={3}' \
        --env='DATALAB_INITIAL_USER_SETTINGS={4}' \
        {0}
+    ExecStop=-/usr/bin/docker stop datalab
+    ExecStopPost=-/usr/bin/docker rm -f datalab
     Restart=always
     RestartSec=1
 
diff --git a/tools/cli/commands/creategpu.py b/tools/cli/commands/creategpu.py
index d3e6c3cba..3d9c269cc 100644
--- a/tools/cli/commands/creategpu.py
+++ b/tools/cli/commands/creategpu.py
@@ -75,7 +75,7 @@
   permissions: 0755
   owner: root
   content: |
-    NVIDIA_DRIVER_VERSION=390.46
+    NVIDIA_DRIVER_VERSION=418.67
     COS_NVIDIA_INSTALLER_CONTAINER=gcr.io/cos-cloud/cos-gpu-installer:latest
     NVIDIA_INSTALL_DIR_HOST=/var/lib/nvidia
     NVIDIA_INSTALL_DIR_CONTAINER=/usr/local/nvidia
@@ -96,8 +96,9 @@
     User=root
     Type=oneshot
     RemainAfterExit=true
+    Environment="HOME=/home/datalab"
     EnvironmentFile=/etc/nvidia-installer-env
-    ExecStartPre=docker-credential-gcr configure-docker
+    ExecStartPre=/usr/bin/docker-credential-gcr configure-docker
     ExecStartPre=/bin/bash -c 'mkdir -p "${{NVIDIA_INSTALL_DIR_HOST}}" && \
         mount --bind "${{NVIDIA_INSTALL_DIR_HOST}}" \
         "${{NVIDIA_INSTALL_DIR_HOST}}" && \
@@ -124,8 +125,9 @@
 
     [Service]
     Environment="HOME=/home/datalab"
-    ExecStartPre=docker-credential-gcr configure-docker
-    ExecStart=/usr/bin/docker run --restart always \
+    ExecStartPre=/usr/bin/docker-credential-gcr configure-docker
+    ExecStart=/usr/bin/docker run \
+       --name=datalab \
        -p '127.0.0.1:8080:8080' \
        -v /mnt/disks/datalab-pd/content:/content \
        -v /mnt/disks/datalab-pd/tmp:/tmp \
@@ -144,6 +146,8 @@
        --env='DATALAB_GIT_AUTHOR={3}' \
        --env='DATALAB_INITIAL_USER_SETTINGS={4}' \
        {0} -c /datalab/run.sh
+    ExecStop=-/usr/bin/docker stop datalab
+    ExecStopPost=-/usr/bin/docker rm -f datalab
     Restart=always
     RestartSec=1