diff --git a/library/psa_crypto.c b/library/psa_crypto.c
index b2e863e6f..875252803 100644
--- a/library/psa_crypto.c
+++ b/library/psa_crypto.c
@@ -1538,40 +1538,32 @@ static psa_status_t psa_finish_key_creation(
 #if defined(MBEDTLS_PSA_CRYPTO_STORAGE_C)
     if( slot->lifetime != PSA_KEY_LIFETIME_VOLATILE )
     {
-        uint8_t *buffer = NULL;
-        size_t buffer_size = 0;
-        size_t length = 0;
+        psa_key_attributes_t attributes = PSA_KEY_ATTRIBUTES_INIT;
+        psa_get_key_slot_attributes( slot, &attributes );
 
 #if defined(MBEDTLS_PSA_CRYPTO_SE_C)
         if( driver != NULL )
         {
-            buffer = (uint8_t*) &slot->data.se.slot_number;
-            length = sizeof( slot->data.se.slot_number );
+            status = psa_save_persistent_key( &attributes,
+                                              (uint8_t*) &slot->data.se,
+                                              sizeof( slot->data.se ) );
         }
         else
 #endif /* MBEDTLS_PSA_CRYPTO_SE_C */
         {
-            buffer_size = PSA_KEY_EXPORT_MAX_SIZE( slot->type,
-                                                   psa_get_key_slot_bits( slot ) );
-            buffer = mbedtls_calloc( 1, buffer_size );
+            size_t buffer_size =
+                PSA_KEY_EXPORT_MAX_SIZE( slot->type,
+                                         psa_get_key_bits( &attributes ) );
+            uint8_t *buffer = mbedtls_calloc( 1, buffer_size );
+            size_t length = 0;
             if( buffer == NULL && buffer_size != 0 )
                 return( PSA_ERROR_INSUFFICIENT_MEMORY );
             status = psa_internal_export_key( slot,
                                               buffer, buffer_size, &length,
                                               0 );
-        }
+            if( status == PSA_SUCCESS )
+                status = psa_save_persistent_key( &attributes, buffer, length );
 
-        if( status == PSA_SUCCESS )
-        {
-            psa_key_attributes_t attributes = PSA_KEY_ATTRIBUTES_INIT;
-            psa_get_key_slot_attributes( slot, &attributes );
-            status = psa_save_persistent_key( &attributes, buffer, length );
-        }
-
-#if defined(MBEDTLS_PSA_CRYPTO_SE_C)
-        if( driver == NULL )
-#endif /* MBEDTLS_PSA_CRYPTO_SE_C */
-        {
             if( buffer_size != 0 )
                 mbedtls_platform_zeroize( buffer, buffer_size );
             mbedtls_free( buffer );
diff --git a/library/psa_crypto_slot_management.c b/library/psa_crypto_slot_management.c
index 6b87ea0b0..e63dcdae6 100644
--- a/library/psa_crypto_slot_management.c
+++ b/library/psa_crypto_slot_management.c
@@ -138,13 +138,12 @@ static psa_status_t psa_load_persistent_key_into_slot( psa_key_slot_t *p_slot )
 #if defined(MBEDTLS_PSA_CRYPTO_SE_C)
     if( psa_key_lifetime_is_external( p_slot->lifetime ) )
     {
-        if( key_data_length != sizeof( p_slot->data.se.slot_number ) )
+        if( key_data_length != sizeof( p_slot->data.se ) )
         {
             status = PSA_ERROR_STORAGE_FAILURE;
             goto exit;
         }
-        memcpy( &p_slot->data.se.slot_number, key_data,
-                sizeof( p_slot->data.se.slot_number ) );
+        memcpy( &p_slot->data.se, key_data, sizeof( p_slot->data.se ) );
     }
     else
 #endif /* MBEDTLS_PSA_CRYPTO_SE_C */