diff --git a/ports/espressif/supervisor/internal_flash.c b/ports/espressif/supervisor/internal_flash.c
index 4c1400469b..3c136cc9fa 100644
--- a/ports/espressif/supervisor/internal_flash.c
+++ b/ports/espressif/supervisor/internal_flash.c
@@ -37,7 +37,8 @@
 #include "py/runtime.h"
 #include "lib/oofatfs/ff.h"
 
-#include "components/spi_flash/include/esp_partition.h"
+#include "esp_ota_ops.h"
+#include "esp_partition.h"
 
 #include "supervisor/flash.h"
 #include "supervisor/usb.h"
@@ -60,9 +61,7 @@ void supervisor_flash_init(void) {
         ESP_PARTITION_SUBTYPE_DATA_FAT,
         NULL);
     #if CIRCUITPY_STORAGE_EXTEND
-    _partition[1] = esp_partition_find_first(ESP_PARTITION_TYPE_APP,
-        ESP_PARTITION_SUBTYPE_APP_OTA_1,
-        NULL);
+    _partition[1] = esp_ota_get_next_update_partition(NULL);
     #endif
 }