diff --git a/include/psa/crypto.h b/include/psa/crypto.h
index 0ba34929b..9f0b13541 100644
--- a/include/psa/crypto.h
+++ b/include/psa/crypto.h
@@ -2724,7 +2724,7 @@ psa_status_t psa_generator_abort(psa_crypto_generator_t *generator);
  * \retval #PSA_ERROR_TAMPERING_DETECTED
  */
 psa_status_t psa_key_derivation(psa_crypto_generator_t *generator,
-                                psa_key_type_t key,
+                                psa_key_slot_t key,
                                 psa_algorithm_t alg,
                                 const uint8_t *salt,
                                 size_t salt_length,
diff --git a/library/psa_crypto.c b/library/psa_crypto.c
index a62fac7b1..210fa5ff4 100644
--- a/library/psa_crypto.c
+++ b/library/psa_crypto.c
@@ -3257,7 +3257,7 @@ static psa_status_t psa_generator_hkdf_setup( psa_hkdf_generator_t *hkdf,
 }
 
 psa_status_t psa_key_derivation( psa_crypto_generator_t *generator,
-                                 psa_key_type_t key,
+                                 psa_key_slot_t key,
                                  psa_algorithm_t alg,
                                  const uint8_t *salt,
                                  size_t salt_length,