diff --git a/hal/stm32_tz.c b/hal/stm32_tz.c
index ef9a5505..86145a15 100644
--- a/hal/stm32_tz.c
+++ b/hal/stm32_tz.c
@@ -105,17 +105,23 @@ void hal_tz_claim_nonsecure_area(uint32_t address, int len)
     while (address < end) {
         if (address < FLASH_BANK2_BASE) {
             page_n = (address - ARCH_FLASH_OFFSET) / FLASH_PAGE_SIZE;
-            bank = 1;
+            bank = 0;
         } else {
             page_n = (address - FLASH_BANK2_BASE) / FLASH_PAGE_SIZE;
-            bank = 2;
+            bank = 1;
         }
+
+#ifdef PLATFORM_stm32h5
+        /* Take into account current swap configuration */
+        if ((FLASH_OPTSR_CUR & FLASH_OPTSR_SWAP_BANK) >> 31)
+            bank = !bank;
+#endif
         reg_idx = page_n / 32;
         pos = page_n % 32;
         hal_flash_wait_complete(bank);
         hal_flash_clear_errors(bank);
         hal_flash_nonsecure_unlock();
-        if (bank == 1)
+        if (bank == 0)
             FLASH_SECBB1[reg_idx] |= ( 1 << pos);
         else
             FLASH_SECBB2[reg_idx] |= ( 1 << pos);
@@ -128,7 +134,7 @@ void hal_tz_claim_nonsecure_area(uint32_t address, int len)
         FLASH_CR = reg | ((page_n << FLASH_CR_PNB_SHIFT) | FLASH_CR_PER);
 #else
         reg = FLASH_CR & (~((FLASH_CR_PNB_MASK << FLASH_CR_PNB_SHIFT) | FLASH_CR_SER | FLASH_CR_BER | FLASH_CR_PG | FLASH_CR_MER));
-        FLASH_CR = reg | ((page_n << FLASH_CR_PNB_SHIFT) | FLASH_CR_SER);
+        FLASH_CR = reg | ((page_n << FLASH_CR_PNB_SHIFT) | FLASH_CR_SER | (bank << 31));
 #endif
 
         DMB();
diff --git a/hal/stm32h5.c b/hal/stm32h5.c
index f69c3488..dbe11783 100644
--- a/hal/stm32h5.c
+++ b/hal/stm32h5.c
@@ -28,6 +28,25 @@
 
 #define PLL_SRC_HSE 1
 
+#if TZ_SECURE()
+static int is_flash_nonsecure(uint32_t address)
+{
+    uint32_t in_bank_offset = address & 0x000FFFFF;
+#ifdef DUALBANK_SWAP
+    if (in_bank_offset >= (WOLFBOOT_PARTITION_BOOT_ADDRESS - FLASHMEM_ADDRESS_SPACE))
+        return 1;
+    else
+        return 0;
+#else
+    if (address >= WOLFBOOT_PARTITION_BOOT_ADDRESS)
+        return 1;
+    else
+        return 0;
+#endif
+}
+#endif
+
+
 static void RAMFUNCTION flash_set_waitstates(unsigned int waitstates)
 {
     uint32_t reg = FLASH_ACR;
@@ -84,11 +103,13 @@ int RAMFUNCTION hal_flash_write(uint32_t address, const uint8_t *data, int len)
     dst = (uint32_t *)address;
 
 #if (TZ_SECURE())
-    if ( ((address < FLASH_BANK2_BASE) && (address >= WOLFBOOT_PARTITION_BOOT_ADDRESS)) ||
-        (address >= WOLFBOOT_PARTITION_UPDATE_ADDRESS))
+    if (is_flash_nonsecure(address)) {
         hal_tz_claim_nonsecure_area(address, len);
+    }
     /* Convert into secure address space */
-    dst = (uint32_t *)((address & (~FLASHMEM_ADDRESS_SPACE)) | FLASH_SECURE_MMAP_BASE);
+    if (((uint32_t)dst & 0x0F000000) == 0x08000000) {
+        dst = (uint32_t *)((address & (~FLASHMEM_ADDRESS_SPACE)) | FLASH_SECURE_MMAP_BASE);
+    }
 #endif
 
     while (i < len) {
@@ -153,6 +174,7 @@ void RAMFUNCTION hal_flash_opt_lock(void)
 }
 
 
+
 int RAMFUNCTION hal_flash_erase(uint32_t address, int len)
 {
     uint32_t end_address;
@@ -173,18 +195,25 @@ int RAMFUNCTION hal_flash_erase(uint32_t address, int len)
         base = FLASHMEM_ADDRESS_SPACE;
         reg = FLASH_CR & (~((FLASH_CR_PNB_MASK << FLASH_CR_PNB_SHIFT) | FLASH_CR_BER));
 
+#if TZ_SECURE()
+        /* When in secure mode, skip erasing non-secure pages: will be erased upon claim */
+        if (is_flash_nonsecure(address)) {
+            return 0;
+        }
+#endif
         if(p >= (FLASH_BANK2_BASE) && (p <= (FLASH_TOP) ))
         {
-#if TZ_SECURE()
-            /* When in secure mode, skip erasing non-secure pages: will be erased upon claim */
-            return 0;
-#endif
             base = FLASH_BANK2_BASE;
             bnksel = 1;
         } else {
             FLASH_CR &= ~FLASH_CR_SER ;
             return 0; /* Address out of range */
         }
+
+        /* Check for swapped banks to invert bnksel */
+        if ((FLASH_OPTSR_CUR & FLASH_OPTSR_SWAP_BANK) >> 31)
+            bnksel = !bnksel;
+
         reg |= ((((p - base)  >> 13) << FLASH_CR_PNB_SHIFT) | FLASH_CR_SER | (bnksel << 31));
         FLASH_CR = reg;
         DMB();
diff --git a/hal/stm32h5.h b/hal/stm32h5.h
index c0176e41..3e901006 100644
--- a/hal/stm32h5.h
+++ b/hal/stm32h5.h
@@ -258,7 +258,7 @@
 
 
 
-#if defined(DUALBANK_SWAP) && defined (__WOLFBOOT)
+#if defined(DUALBANK_SWAP)
     #define FLASH_OPTSR_CUR   (*(volatile uint32_t *)(FLASH_BASE + 0x50))
     #define FLASH_OPTSR_PRG   (*(volatile uint32_t *)(FLASH_BASE + 0x54))
     #define FLASH_OPTSR_SWAP_BANK (1 << 31)
diff --git a/test-app/app_stm32h5.c b/test-app/app_stm32h5.c
index 5d80849c..e60d055a 100644
--- a/test-app/app_stm32h5.c
+++ b/test-app/app_stm32h5.c
@@ -279,7 +279,7 @@ static int cmd_update_xmodem(const char *args)
     printf("Erasing update partition...");
     fflush(stdout);
     hal_flash_unlock();
-    //hal_flash_erase(dst_flash, WOLFBOOT_PARTITION_SIZE);
+    hal_flash_erase(dst_flash, WOLFBOOT_PARTITION_SIZE);
     printf("Done.\r\n");
 
     printf("Waiting for XMODEM transfer...\r\n");
@@ -400,6 +400,7 @@ static int cmd_info(const char *args)
     printf("\r\n");
     printf("System information\r\n");
     printf("====================================\r\n");
+    printf("Flash banks are %sswapped.\r\n", ((FLASH_OPTSR_CUR & (FLASH_OPTSR_SWAP_BANK)) == 0)?"not ":"");
     printf("Firmware version : 0x%lx\r\n", wolfBoot_current_firmware_version());
     if (update_fw_version != 0) {
         printf("Candidate firmware version : 0x%lx\r\n", update_fw_version);