diff --git a/library/psa_crypto_pake.h b/library/psa_crypto_pake.h
index c7bf270..9256f5a 100644
--- a/library/psa_crypto_pake.h
+++ b/library/psa_crypto_pake.h
@@ -442,7 +442,7 @@
  */
 psa_status_t mbedtls_psa_pake_get_implicit_key(
     mbedtls_psa_pake_operation_t *operation,
-    psa_key_derivation_operation_t *output);
+    uint8_t *output, size_t *output_size);
 
 /** Abort a PAKE operation.
  *