diff --git a/rebuild.sh b/rebuild.sh
index 00d1700..5996f0b 100755
--- a/rebuild.sh
+++ b/rebuild.sh
@@ -7,10 +7,12 @@ git diff -U0 *.nix
 
 # add generation comment to flake.nix
 gen=$(nixos-rebuild list-generations | grep current | awk '{print $1}')
-if grep -q '^# generation:' flake.nix; then
-    sed -i "1s/^# generation:.*/# generation: $gen/" flake.nix
+if sed -n '3p' flake.nix | grep -q '^# generation:'; then
+    # replace the comment on line 3
+    sed -i "3s/^# generation:.*/# generation: $gen/" flake.nix
 else
-    sed -i "1i# generation: $gen" flake.nix
+    # insert comment on line 3
+    sed -i "3i# generation: $gen" flake.nix
 fi
 
 echo "nixos rebuilding..."