diff --git a/include/psa/crypto_struct.h b/include/psa/crypto_struct.h
index bf178ec..6a018e1 100644
--- a/include/psa/crypto_struct.h
+++ b/include/psa/crypto_struct.h
@@ -374,9 +374,17 @@
 static inline void psa_set_key_id( psa_key_attributes_t *attributes,
                                    mbedtls_svc_key_id_t key )
 {
+    psa_key_lifetime_t lifetime = attributes->core.lifetime;
+
     attributes->core.id = key;
-    if( attributes->core.lifetime == PSA_KEY_LIFETIME_VOLATILE )
-        attributes->core.lifetime = PSA_KEY_LIFETIME_PERSISTENT;
+
+    if( PSA_KEY_LIFETIME_IS_VOLATILE( lifetime ) )
+    {
+        attributes->core.lifetime =
+            PSA_KEY_LIFETIME_FROM_PERSISTENCE_AND_LOCATION(
+                PSA_KEY_LIFETIME_PERSISTENT,
+                PSA_KEY_LIFETIME_GET_LOCATION( lifetime ) );
+    }
 }
 
 static inline mbedtls_svc_key_id_t psa_get_key_id(
@@ -397,7 +405,7 @@
                                         psa_key_lifetime_t lifetime)
 {
     attributes->core.lifetime = lifetime;
-    if( lifetime == PSA_KEY_LIFETIME_VOLATILE )
+    if( PSA_KEY_LIFETIME_IS_VOLATILE( lifetime ) )
     {
 #ifdef MBEDTLS_PSA_CRYPTO_KEY_ID_ENCODES_OWNER
         attributes->core.id.key_id = 0;