diff --git a/scripts/Windows/OS/Install.ps1 b/scripts/Windows/OS/Install.ps1
index 19e947bac..9f2fd3cf5 100644
--- a/scripts/Windows/OS/Install.ps1
+++ b/scripts/Windows/OS/Install.ps1
@@ -69,7 +69,7 @@ function Start-WindowsInstallation([Context] $context) {
 
                 winget install --accept-source-agreements --accept-package-agreements -e --id AutoHotkey.AutoHotkey;
 
-                $context.Set("SoftwarePrerequisitesInstalled", 1, "DWord");
+                $context.SetStage($driverStage);
                 break;
             }
             $driverStage {
@@ -81,7 +81,7 @@ function Start-WindowsInstallation([Context] $context) {
                 }
 
                 Write-Information "Finished installing drivers";
-                $context.SetStage("DriversInstalled", 1, "DWord");
+                $context.SetStage($softwareStage);
                 $context.Reboot();
                 exit;
             }