diff --git a/load_balancer.tf b/load_balancer.tf
index 5408677..316980f 100644
--- a/load_balancer.tf
+++ b/load_balancer.tf
@@ -42,6 +42,7 @@ resource "azurerm_lb_rule" "monitoring_vxlan_lb_rule" {
   backend_address_pool_ids = [
     azurerm_lb_backend_address_pool.monitoring_pool.id
   ]
+  probe_id = azurerm_lb_probe.sensor_health_check_probe.id
 }
 
 resource "azurerm_lb_rule" "monitoring_geneve_lb_rule" {
@@ -54,6 +55,7 @@ resource "azurerm_lb_rule" "monitoring_geneve_lb_rule" {
   backend_address_pool_ids = [
     azurerm_lb_backend_address_pool.monitoring_pool.id
   ]
+  probe_id = azurerm_lb_probe.sensor_health_check_probe.id
 }
 
 resource "azurerm_lb_rule" "monitoring_health_check_rule" {
@@ -79,4 +81,5 @@ resource "azurerm_lb_rule" "management_lb_rule" {
   backend_address_pool_ids = [
     azurerm_lb_backend_address_pool.management_pool.id
   ]
+  probe_id = azurerm_lb_probe.sensor_health_check_probe.id
 }
diff --git a/nat_gateway.tf b/nat_gateway.tf
deleted file mode 100644
index 72d1bc2..0000000
--- a/nat_gateway.tf
+++ /dev/null
@@ -1,27 +0,0 @@
-resource "azurerm_public_ip" "nat_gw_ip" {
-  name                = var.nat_gateway_ip_name
-  location            = var.location
-  resource_group_name = var.resource_group_name
-  allocation_method   = "Static"
-  sku                 = "Standard"
-
-  tags = var.tags
-}
-
-resource "azurerm_nat_gateway" "lb_nat_gw" {
-  name                = var.nat_gateway_name
-  location            = var.location
-  resource_group_name = var.resource_group_name
-
-  tags = var.tags
-}
-
-resource "azurerm_subnet_nat_gateway_association" "nat_gw_association" {
-  subnet_id      = azurerm_subnet.subnet.id
-  nat_gateway_id = azurerm_nat_gateway.lb_nat_gw.id
-}
-
-resource "azurerm_nat_gateway_public_ip_association" "public_ip_association" {
-  nat_gateway_id       = azurerm_nat_gateway.lb_nat_gw.id
-  public_ip_address_id = azurerm_public_ip.nat_gw_ip.id
-}
\ No newline at end of file
diff --git a/outputs.tf b/outputs.tf
index 3ab2492..40e42be 100644
--- a/outputs.tf
+++ b/outputs.tf
@@ -2,14 +2,6 @@ output "internal_load_balancer_name" {
   value = azurerm_lb.scale_set_lb.name
 }
 
-output "nat_gateway_public_ip_name" {
-  value = azurerm_public_ip.nat_gw_ip.name
-}
-
-output "nat_gateway_name" {
-  value = azurerm_nat_gateway.lb_nat_gw.name
-}
-
 output "sensor_identity_principal_id" {
   value = azurerm_linux_virtual_machine_scale_set.sensor_scale_set.identity[0].principal_id
 }
diff --git a/scale_set.tf b/scale_set.tf
index 6948fb6..d046e2f 100644
--- a/scale_set.tf
+++ b/scale_set.tf
@@ -24,6 +24,8 @@ resource "azurerm_linux_virtual_machine_scale_set" "sensor_scale_set" {
   }
 
   health_probe_id = azurerm_lb_probe.sensor_health_check_probe.id
+  upgrade_mode    = "Automatic"
+
   network_interface {
     name    = "management-nic"
     primary = true
@@ -51,6 +53,10 @@ resource "azurerm_linux_virtual_machine_scale_set" "sensor_scale_set" {
   }
 
   tags = var.tags
+
+  depends_on = [
+    azurerm_lb_rule.monitoring_health_check_rule,
+  ]
 }
 
 resource "azurerm_monitor_autoscale_setting" "auto_scale_config" {
diff --git a/variables.tf b/variables.tf
index 473fcf5..5ee355e 100644
--- a/variables.tf
+++ b/variables.tf
@@ -57,18 +57,6 @@ variable "sensor_admin_username" {
   default     = "corelight"
 }
 
-variable "nat_gateway_ip_name" {
-  description = "The resource name of the VMSS NAT Gateway public IP resource"
-  type        = string
-  default     = "cl-nat-gw-ip"
-}
-
-variable "nat_gateway_name" {
-  description = "The resource name of the VMSS NAT Gateway resource"
-  type        = string
-  default     = "cl-sensor-nat-gw"
-}
-
 variable "autoscale_setting_name" {
   description = "The VMSS autoscale monitor name"
   type        = string