diff --git a/bin/git-check-assertions b/bin/git-check-assertions
index 53fc277..fec0805 100755
--- a/bin/git-check-assertions
+++ b/bin/git-check-assertions
@@ -72,6 +72,9 @@ else
   echo "No main or master branch found." >&2
   exit 1
 fi
+restore() {
+  git restore -q --staged --worktree .
+}
 base="$(git merge-base "$base_branch" HEAD)"
 mapfile -t commits < <(git rev-list --reverse "${base}..HEAD")
 
@@ -92,12 +95,12 @@ for commit_hash in "${commits[@]}"; do
     printf '%s\n' "$block" | sed 's/^/> /'
     if ! bash -euo pipefail -c "$block"; then
       echo "git-check-assertions block failed in $commit_hash" >&2
-      git restore -q .
+      restore
       echo "Returning to $orig_ref"
       git checkout -q "$orig_ref"
       exit 1
     fi
-    git restore -q .
+    restore
   fi
   echo
 done