diff --git a/library/psa_crypto.c b/library/psa_crypto.c
index 8c59cf6..663b645 100644
--- a/library/psa_crypto.c
+++ b/library/psa_crypto.c
@@ -5174,7 +5174,8 @@
     {
         return( PSA_ERROR_NOT_SUPPORTED );
     }
-#if defined(MBEDTLS_PSA_BUILTIN_ALG_HKDF_EXTRACT)
+#if defined(MBEDTLS_PSA_BUILTIN_ALG_HKDF_EXTRACT) || \
+    defined(MBEDTLS_PSA_BUILTIN_ALG_TLS12_PSK_TO_MS)
     if( PSA_ALG_IS_HKDF_EXTRACT( kdf_alg ) ||
         ( kdf_alg == PSA_ALG_TLS12_ECJPAKE_TO_PMS ) )
         operation->capacity = hash_size;