diff --git a/locals.tf b/locals.tf
index 050fef8..c07017b 100644
--- a/locals.tf
+++ b/locals.tf
@@ -1,6 +1,6 @@
 locals {
   public_ip = {
-    name = coalesce(var.name_overrides.public_ip, "pip-${var.stage}-${var.virtual_machine_config.hostname}-01-${var.virtual_machine_config.location}")
+    name = coalesce(var.name_overrides.public_ip, "pip-${var.public_ip_config.stage}-${var.virtual_machine_config.hostname}-01-${var.virtual_machine_config.location}")
   }
 
   nic = {
@@ -13,5 +13,5 @@ locals {
     tags = merge(var.tags, { "Severity Group Monthly" = var.virtual_machine_config.severity_group,  "Update allowed" = local.update_allowed })
   }
   os_disk_name   = coalesce(var.name_overrides.os_disk, "disk-${var.virtual_machine_config.hostname}-Os")
-  update_allowed = var.update_allowed ? "yes" : "no"
+  update_allowed = var.virtual_machine_config.update_allowed ? "yes" : "no"
 }
\ No newline at end of file
diff --git a/variables.tf b/variables.tf
index b17b681..8c8eefa 100644
--- a/variables.tf
+++ b/variables.tf
@@ -2,6 +2,7 @@ variable "public_ip_config" {
   type = object({
     enabled           = bool
     allocation_method = optional(string, "Static")
+    stage             = string
   })
   default = {
     enabled = false
@@ -14,7 +15,8 @@ variable "public_ip_config" {
   description = <<-DOC
   ```
     enabled: Optionally select true if a public ip should be created. Defaults to false.
-    allocation_method: The allocation method of the public ip that will be created. Defaults to static.      
+    allocation_method: The allocation method of the public ip that will be created. Defaults to static.
+    stage: The stage of this PIP. Ex: prd, dev, tst, ...
   ```
   DOC
 }
@@ -101,6 +103,7 @@ variable "virtual_machine_config" {
     availability_set_id               = optional(string)
     proximity_placement_group_id      = optional(string)
     severity_group                    = string
+    update_allowed                    = optional(bool, true)
   })
   validation {
     condition     = contains(["None", "ReadOnly", "ReadWrite"], var.virtual_machine_config.os_disk_caching)
@@ -111,7 +114,13 @@ variable "virtual_machine_config" {
     error_message = "Possible values are Standard_LRS, StandardSSD_LRS, Premium_LRS, StandardSSD_ZRS and Premium_ZRS for os_disk_storage_type."
   }
   validation {
-    condition     = (contains(["Premium_LRS", "Premium_ZRS"], var.virtual_machine_config.os_disk_storage_type) && var.virtual_machine_config.os_disk_write_accelerator_enabled == true && var.virtual_machine_config.os_disk_caching == "None") || (var.virtual_machine_config.os_disk_write_accelerator_enabled == false)
+    condition = (
+        var.virtual_machine_config.os_disk_write_accelerator_enabled == true && 
+        contains(["Premium_LRS", "Premium_ZRS"], var.virtual_machine_config.os_disk_storage_type) && 
+        contains(["None", "ReadOnly"], var.virtual_machine_config.os_disk_caching)
+      ) || (
+        var.virtual_machine_config.os_disk_write_accelerator_enabled == false
+      )
     error_message = "os_disk_write_accelerator_enabled can only be activated on Premium disks and caching deactivated."
   }
   validation {
@@ -137,17 +146,11 @@ variable "virtual_machine_config" {
       be activated on Premium disks and caching deactivated. Defaults to false.
     proximity_placement_group_id: (Optional) The ID of the Proximity Placement Group which the Virtual Machine should be assigned to.
     severity_group: (Required) Sets tag 'Severity Group Monthly' to a specific time and date when an update will be done automatically.
+    update_allowed: Sets tag 'Update allowed' to yes or no to specify if this VM should currently receive updates.
   ```
   DOC
 }
 
-variable "update_allowed" {
-  type        = bool
-  default     = true
-  nullable    = false
-  description = "Set the tag `Update allowed`. `True` will set `yes`, `false` to `no`."
-}
-
 variable "data_disks" {
   type = map(object({
     lun                        = number
@@ -229,9 +232,3 @@ variable "tags" {
   nullable    = false
   description = "A map of tags that will be set on every resource this module creates."
 }
-
-variable "stage" {
-  type        = string
-  nullable    = false
-  description = "The stage of this VM like prd, dev, tst, ..."
-}