diff --git a/library/psa_crypto.c b/library/psa_crypto.c
index 0b21eb083..553834f76 100644
--- a/library/psa_crypto.c
+++ b/library/psa_crypto.c
@@ -5175,12 +5175,13 @@ static psa_status_t psa_key_derivation_setup_kdf(
         return( PSA_ERROR_NOT_SUPPORTED );
     }
 #if defined(MBEDTLS_PSA_BUILTIN_ALG_HKDF_EXTRACT) || \
-    defined(MBEDTLS_PSA_BUILTIN_ALG_TLS12_PSK_TO_MS)
+    defined(MBEDTLS_PSA_BUILTIN_ALG_TLS12_ECJPAKE_TO_PMS)
     if( PSA_ALG_IS_HKDF_EXTRACT( kdf_alg ) ||
         ( kdf_alg == PSA_ALG_TLS12_ECJPAKE_TO_PMS ) )
         operation->capacity = hash_size;
     else
-#endif /* MBEDTLS_PSA_BUILTIN_ALG_HKDF_EXTRACT */
+#endif /* MBEDTLS_PSA_BUILTIN_ALG_HKDF_EXTRACT ||
+          MBEDTLS_PSA_BUILTIN_ALG_TLS12_ECJPAKE_TO_PMS */
         operation->capacity = 255 * hash_size;
     return( PSA_SUCCESS );
 }