diff --git a/locals.tf b/locals.tf
index 5e4745b..c0a66f2 100644
--- a/locals.tf
+++ b/locals.tf
@@ -2,4 +2,9 @@ locals {
   sso_permission_sets = var.sso_permission_sets
   organization_config = var.organization_config
   enable_sso          = var.enable_sso
+  accounts = flatten([
+    for unit_name, unit in local.organization_config["units"] : [
+      for account_name in keys(local.organization_config["units"][unit_name]["accounts"]) : local.organization_config["units"][unit_name]["accounts"][account_name]
+    ]
+  ])
 }
diff --git a/sso.tf b/sso.tf
index 8493d2d..2b5102c 100644
--- a/sso.tf
+++ b/sso.tf
@@ -3,11 +3,7 @@ data "aws_ssoadmin_instances" "ssoadmin_instances" {}
 data "aws_identitystore_group" "aws" {
   for_each = local.enable_sso ? toset(
     flatten([
-      for account in flatten([
-        for unit_name, unit in local.organization_config["units"] : [
-          for account_name in keys(local.organization_config["units"][unit_name]["accounts"]) : local.organization_config["units"][unit_name]["accounts"][account_name]
-        ]
-      ]) : keys(lookup(account, "group_assignments", {}))
+      for account in local.accounts : keys(lookup(account, "group_assignments", {}))
     ])
   ) : toset([])
 
@@ -24,11 +20,7 @@ data "aws_identitystore_group" "aws" {
 data "aws_identitystore_user" "aws" {
   for_each = local.enable_sso ? toset(
     flatten([
-      for account in flatten([
-        for unit_name, unit in local.organization_config["units"] : [
-          for account_name in keys(local.organization_config["units"][unit_name]["accounts"]) : local.organization_config["units"][unit_name]["accounts"][account_name]
-        ]
-      ]) : keys(lookup(account, "user_assignments", {}))
+      for account in local.accounts : keys(lookup(account, "user_assignments", {}))
     ])
   ) : toset([])
 
@@ -53,16 +45,14 @@ resource "aws_ssoadmin_permission_set" "permission_set" {
 }
 
 resource "aws_ssoadmin_managed_policy_attachment" "attachment" {
-  for_each = local.enable_sso ? {
-    for attachment in flatten([
-      for permission_set_name, permission_set in local.sso_permission_sets : {
-        for managed_policy_name in lookup(permission_set, "managed_policies", []) : "${permission_set_name}_${managed_policy_name}" => {
-          permission_set_name = permission_set_name
-          managed_policy_name = managed_policy_name
-        }
+  for_each = local.enable_sso ? merge([
+    for permission_set_name, permission_set in local.sso_permission_sets : {
+      for managed_policy_name in permission_set["managed_policies"] : "${permission_set_name}_${managed_policy_name}" => {
+        permission_set_name = permission_set_name
+        managed_policy_name = managed_policy_name
       }
-    ]) : keys(attachment)[0] => attachment[keys(attachment)[0]]
-  } : {}
+    }
+  ]...) : {}
 
   instance_arn       = tolist(data.aws_ssoadmin_instances.ssoadmin_instances.arns)[0]
   managed_policy_arn = "arn:aws:iam::aws:policy/${each.value["managed_policy_name"]}"
@@ -82,21 +72,19 @@ resource "aws_ssoadmin_permission_set_inline_policy" "policy" {
 }
 
 resource "aws_ssoadmin_account_assignment" "group_assignment" {
-  for_each = local.enable_sso ? {
-    for assignment in flatten([
-      for unit_name, unit in local.organization_config["units"] : [
-        for account_name in keys(local.organization_config["units"][unit_name]["accounts"]) : [
-          for group_name, group_assignments in lookup(local.organization_config["units"][unit_name]["accounts"][account_name], "group_assignments", {}) : {
-            for permission_set in group_assignments["permission_sets"] : "${account_name}_${group_name}_${permission_set}" => {
-              account_name   = account_name
-              group_name     = group_name
-              permission_set = permission_set
-            }
+  for_each = local.enable_sso ? merge(flatten([
+    for unit_name, unit in local.organization_config["units"] : [
+      for account_name in keys(local.organization_config["units"][unit_name]["accounts"]) : [
+        for group_name, group_assignments in lookup(local.organization_config["units"][unit_name]["accounts"][account_name], "group_assignments", {}) : {
+          for permission_set in group_assignments["permission_sets"] : "${account_name}_${group_name}_${permission_set}" => {
+            account_name   = account_name
+            group_name     = group_name
+            permission_set = permission_set
           }
-        ]
+        }
       ]
-    ]) : keys(assignment)[0] => assignment[keys(assignment)[0]]
-  } : {}
+    ]
+  ])...) : {}
 
   instance_arn       = aws_ssoadmin_permission_set.permission_set[each.value["permission_set"]].instance_arn
   permission_set_arn = aws_ssoadmin_permission_set.permission_set[each.value["permission_set"]].arn
@@ -109,21 +97,19 @@ resource "aws_ssoadmin_account_assignment" "group_assignment" {
 }
 
 resource "aws_ssoadmin_account_assignment" "user_assignment" {
-  for_each = local.enable_sso ? {
-    for assignment in flatten([
-      for unit_name, unit in local.organization_config["units"] : [
-        for account_name in keys(local.organization_config["units"][unit_name]["accounts"]) : [
-          for user_name, user_assignments in lookup(local.organization_config["units"][unit_name]["accounts"][account_name], "user_assignments", {}) : {
-            for permission_set in user_assignments["permission_sets"] : "${account_name}_${user_name}_${permission_set}" => {
-              account_name   = account_name
-              user_name      = user_name
-              permission_set = permission_set
-            }
+  for_each = local.enable_sso ? merge(flatten([
+    for unit_name, unit in local.organization_config["units"] : [
+      for account_name in keys(local.organization_config["units"][unit_name]["accounts"]) : [
+        for user_name, user_assignments in lookup(local.organization_config["units"][unit_name]["accounts"][account_name], "user_assignments", {}) : {
+          for permission_set in user_assignments["permission_sets"] : "${account_name}_${user_name}_${permission_set}" => {
+            account_name   = account_name
+            user_name      = user_name
+            permission_set = permission_set
           }
-        ]
+        }
       ]
-    ]) : keys(assignment)[0] => assignment[keys(assignment)[0]]
-  } : {}
+    ]
+  ])...) : {}
 
   instance_arn       = aws_ssoadmin_permission_set.permission_set[each.value["permission_set"]].instance_arn
   permission_set_arn = aws_ssoadmin_permission_set.permission_set[each.value["permission_set"]].arn
diff --git a/versions.tf b/versions.tf
index b188b75..ee1ff3b 100644
--- a/versions.tf
+++ b/versions.tf
@@ -1,6 +1,9 @@
 terraform {
   required_version = ">= 1.1.5"
   required_providers {
-    aws = "~> 5.0"
+    aws = {
+      source  = "hashicorp/aws"
+      version = ">= 5.0"
+    }
   }
 }