diff --git a/Update-VMGpuPartitionDriver.ps1 b/Update-VMGpuPartitionDriver.ps1
index 9645833..5215258 100644
--- a/Update-VMGpuPartitionDriver.ps1
+++ b/Update-VMGpuPartitionDriver.ps1
@@ -68,34 +68,34 @@ foreach ($d in $drivers) {
 }
 
 $VM = Get-VM -VMName "GPUP"
+$VHD = Get-VHD -VMId $VM.VMId
 
 If ($VM.state -eq "Running") {
-[bool]$state_was_running = $true
-}
+    [bool]$state_was_running = $true
+    }
 
 if ($VM.state -ne "Off"){
-"Attemping to shutdown VM..."
-Stop-VM -Name $VMName -Force
-} 
+    "Attemping to shutdown VM..."
+    Stop-VM -Name $VMName -Force
+    } 
 
 While ($VM.State -ne "Off") {
-Start-Sleep -s 3
-"Waiting for VM to shutdown - make sure there are no unsaved documents..."
-}
-
-$VHD = Get-VHD -VMId $VM.VMId
+    Start-Sleep -s 3
+    "Waiting for VM to shutdown - make sure there are no unsaved documents..."
+    }
 
 "Mounting Drive..."
 $DriveLetter = (Mount-VHD -Path $VHD.Path -PassThru | Get-Disk | Get-Partition | Get-Volume | Where-Object {$_.DriveLetter} | ForEach-Object DriveLetter) + ":"
 
 "Copying GPU Files - this could take a while..."
-
 Add-VMGPUPartitionAdapterFiles -hostname $Hostname -DriveLetter $DriveLetter -GPUName $GPUName
 
 "Dismounting Drive..."
 Dismount-VHD -Path $VHD.Path
 
 If ($state_was_running){
-"Previous State was running so starting VM..."
-Start-VM $VMName
-}
\ No newline at end of file
+    "Previous State was running so starting VM..."
+    Start-VM $VMName
+    }
+
+"Done..."
\ No newline at end of file