diff --git a/library/psa_crypto.c b/library/psa_crypto.c
index 84b10df..93c9ce4 100644
--- a/library/psa_crypto.c
+++ b/library/psa_crypto.c
@@ -1348,6 +1348,18 @@
     }
     slot->type = attributes->type;
 
+#if defined(MBEDTLS_PSA_CRYPTO_SE_C)
+    /* Find a slot number. Don't yet mark it as allocated in case
+     * the key creation fails or there is a power failure. */
+    if( *p_drv != NULL )
+    {
+        status = psa_find_se_slot_for_key( attributes, *p_drv,
+                                           &slot->data.se.slot_number );
+        if( status != PSA_SUCCESS )
+            return( status );
+    }
+#endif /* MBEDTLS_PSA_CRYPTO_SE_C */
+
     return( status );
 }
 
@@ -1405,6 +1417,18 @@
     }
 #endif /* defined(MBEDTLS_PSA_CRYPTO_STORAGE_C) */
 
+#if defined(MBEDTLS_PSA_CRYPTO_SE_C)
+    if( driver != NULL )
+    {
+        status = psa_save_se_persistent_data( driver );
+        if( status != PSA_SUCCESS )
+        {
+            psa_destroy_persistent_key( slot->persistent_storage_id );
+            return( status );
+        }
+    }
+#endif /* MBEDTLS_PSA_CRYPTO_SE_C */
+
     return( status );
 }