diff --git a/informers.go b/informers.go
index e573082..d98ea50 100644
--- a/informers.go
+++ b/informers.go
@@ -1,6 +1,7 @@
 package mass_machine_type_transition
 
 import (
+	"context"
 	"fmt"
 	"strings"
 	"time"
@@ -13,7 +14,11 @@ import (
 	"kubevirt.io/client-go/kubecli"
 )
 
-var needsRestart = 0
+// using this as a const allows us to easily modify the program to update if a newer version is released
+// we generally want to be updating the machine types to the most recent version
+const latestVersion = "rhel9.2.0"
+
+var vmiList []string{}
 
 func getVirtCli() (kubecli.KubevirtClient, error) {
 	clientConfig, err := kubecli.GetKubevirtClientConfig()
@@ -49,9 +54,12 @@ func updateMachineTypes(virtCli kubecli.KubevirtClient) error {
 		machineType := vm.Spec.Template.Spec.Domain.Machine.Type
 		machineTypeSubstrings := strings.Split(machineType, "-")
 		version := machineTypeSubstrings[2]
+		if len(machineTypeSubstrings != 3) {
+			return nil
+		}
 		
 		if strings.Contains(version, "rhel") && version < "rhel9.0.0" {
-			machineTypeSubstrings[2] = "rhel9.0.0"
+			machineTypeSubstrings[2] = latestVersion
 			machineType = strings.Join(machineTypeSubstrings, "-")
 			updateMachineType := fmt.Sprintf(`{"spec": {"template": {"spec": {"domain": {"machine": {"type": "%s"}}}}}}`, machineType)
 			
@@ -62,17 +70,25 @@ func updateMachineTypes(virtCli kubecli.KubevirtClient) error {
 			
 			// add label to running VMs that a restart is required for change to take place
 			if vm.Status.Ready {
-				addWarningLabel(virtCli, &vm)
+				err = addWarningLabel(virtCli, &vm)
+				if err != nil {
+					return err
+				}
 			}
 		}
 	}
 	return nil
 }
 
-func addWarningLabel (virtCli kubecli.KubevirtClient, vm *k6tv1.VirtualMachine) {
+func addWarningLabel (virtCli kubecli.KubevirtClient, vm *k6tv1.VirtualMachine) error {
 	addLabel := fmt.Sprint(`{"metadata": {"labels": {"restart-vm-required": "true"}}}}}`)
-	virtCli.VirtualMachine(vm.Namespace).Patch(vm.Name, types.StrategicMergePatchType, []byte(addLabel), &k8sv1.PatchOptions{})
-	needsRestart++
+	virtCli.VirtualMachineInstance(vm.Namespace).Patch(context.Background(), vm.Name, types.StrategicMergePatchType, []byte(addLabel), &k8sv1.PatchOptions{})
+	if err != nil {
+		return err
+	}
+	vmiList = append(vmiList, vm.Name)
+	
+	return nil
 }
 
 func removeWarningLabel(obj interface{}) {
@@ -86,7 +102,21 @@ func removeWarningLabel(obj interface{}) {
 		return
 	}
 	
-	removeLabel := fmt.Sprint(`[{"op": "remove", "path": "metadata/labels/restart-vm-required"}]`)
-	virtCli.VirtualMachine(vmi.Namespace).Patch(vmi.Name, types.StrategicMergePatchType, []byte(removeLabel), &k8sv1.PatchOptions{})
-	needsRestart--
+	//check if deleted VMI is in list of VMIs that need to be restarted
+	vmiIndex = searchVMIList(vmi.Name)
+	if  vmiIndex == -1 {
+		return
+	}
+	
+	// remove deleted VMI from list
+	vmiList = append(vmiList[:vmiIndex], vmiList[vmiIndex+1:]
+}
+
+func searchVMIList(vmiName string) int {
+	for i, element := range vmiList {
+		if element == vmiName {
+			return i
+		}
+	}
+	return -1
 }
diff --git a/main.go b/main.go
index 07cf945..d25c9e2 100644
--- a/main.go
+++ b/main.go
@@ -18,12 +18,9 @@ func main() {
 	
 	updateMachineTypes(virtCli)
 	
-	// for now this uses the number of vms to be restarted to determine
-	// when to exit the job rather than seeing if the actual label exists
-	// on any VMs before exiting the job
-	
-	for (needsRestart > 0) {
-	}
+	// need to implement probably a chan of some kind to inform the
+	// program of when to exit after waiting for all VMs to be
+	// rebooted
 	
 	os.Exit(0)
 }