diff --git a/library/ssl_tls13_keys.c b/library/ssl_tls13_keys.c
index ff5627a..8ec3570 100644
--- a/library/ssl_tls13_keys.c
+++ b/library/ssl_tls13_keys.c
@@ -1463,7 +1463,7 @@
      */
     if (mbedtls_ssl_tls13_key_exchange_mode_with_ephemeral(ssl)) {
         if (mbedtls_ssl_tls13_named_group_is_ecdhe(handshake->offered_group_id)) {
-#if defined(MBEDTLS_PK_CAN_ECDH)
+#if defined(PSA_WANT_ALG_ECDH)
             /* Compute ECDH shared secret. */
             psa_status_t status = PSA_ERROR_GENERIC_ERROR;
             psa_key_attributes_t key_attributes = PSA_KEY_ATTRIBUTES_INIT;
@@ -1499,7 +1499,7 @@
             }
 
             handshake->ecdh_psa_privkey = MBEDTLS_SVC_KEY_ID_INIT;
-#endif /* MBEDTLS_PK_CAN_ECDH */
+#endif /* PSA_WANT_ALG_ECDH */
         } else {
             MBEDTLS_SSL_DEBUG_MSG(1, ("Group not supported."));
             return MBEDTLS_ERR_SSL_FEATURE_UNAVAILABLE;