diff --git a/library/psa_crypto.c b/library/psa_crypto.c
index cbdc912..981b6f8 100644
--- a/library/psa_crypto.c
+++ b/library/psa_crypto.c
@@ -5175,7 +5175,8 @@
         return( PSA_ERROR_NOT_SUPPORTED );
     }
 #if defined(MBEDTLS_PSA_BUILTIN_ALG_HKDF_EXTRACT)
-    if( PSA_ALG_IS_HKDF_EXTRACT( kdf_alg ) )
+    if( PSA_ALG_IS_HKDF_EXTRACT( kdf_alg ) ||
+        PSA_ALG_IS_TLS12_ECJPAKE_TO_PMS ( kdf_alg ))
         operation->capacity = hash_size;
     else
 #endif /* MBEDTLS_PSA_BUILTIN_ALG_HKDF_EXTRACT */