diff --git a/library/pk.c b/library/pk.c
index b2f6812..331ed6c 100644
--- a/library/pk.c
+++ b/library/pk.c
@@ -146,6 +146,7 @@
 int mbedtls_pk_setup_psa( mbedtls_pk_context *ctx, const psa_key_slot_t key )
 {
     const mbedtls_pk_info_t * const info = &mbedtls_pk_opaque_psa_info;
+    psa_key_slot_t *pk_ctx;
 
     if( ctx == NULL || ctx->pk_info != NULL )
         return( MBEDTLS_ERR_PK_BAD_INPUT_DATA );
@@ -153,11 +154,11 @@
     if( ( ctx->pk_ctx = info->ctx_alloc_func() ) == NULL )
         return( MBEDTLS_ERR_PK_ALLOC_FAILED );
 
-    /* coming soon: remember key */
-    (void) key;
-
     ctx->pk_info = info;
 
+    pk_ctx = (psa_key_slot_t *) ctx->pk_ctx;
+    *pk_ctx = key;
+
     return( 0 );
 }
 #endif /* MBEDTLS_USE_PSA_CRYPTO */