diff --git a/library/ssl_tls13_client.c b/library/ssl_tls13_client.c
index 6aae51b..dac1a1e 100644
--- a/library/ssl_tls13_client.c
+++ b/library/ssl_tls13_client.c
@@ -200,7 +200,7 @@
         ssl->handshake->ecdh_psa_privkey = MBEDTLS_SVC_KEY_ID_INIT;
         return 0;
     } else
-#endif /* MBEDTLS_ECDH_C */
+#endif /* MBEDTLS_PK_CAN_ECDH */
     if (0 /* other KEMs? */) {
         /* Do something */
     }
@@ -237,7 +237,7 @@
 #else
     ((void) ssl);
     ((void) group_id);
-#endif /* MBEDTLS_ECDH_C */
+#endif /* MBEDTLS_PK_CAN_ECDH */
 
     /*
      * Add DHE named groups here.
@@ -326,7 +326,7 @@
         /* Write key_exchange_length */
         MBEDTLS_PUT_UINT16_BE(key_exchange_len, group, 2);
     } else
-#endif /* MBEDTLS_ECDH_C */
+#endif /* MBEDTLS_PK_CAN_ECDH */
     if (0 /* other KEMs? */) {
         /* Do something */
     } else {
@@ -375,7 +375,7 @@
                                              const unsigned char *buf,
                                              const unsigned char *end)
 {
-#if defined(MBEDTLS_ECDH_C)
+#if defined(MBEDTLS_PK_CAN_ECDH)
     const unsigned char *p = buf;
     int selected_group;
     int found = 0;
@@ -480,7 +480,7 @@
         return MBEDTLS_ERR_SSL_HANDSHAKE_FAILURE;
     }
 
-#if defined(MBEDTLS_ECDH_C)
+#if defined(MBEDTLS_PK_CAN_ECDH)
     if (mbedtls_ssl_tls13_named_group_is_ecdhe(group)) {
         if (mbedtls_ssl_get_psa_curve_info_from_tls_id(group, NULL, NULL)
             == PSA_ERROR_NOT_SUPPORTED) {
@@ -496,7 +496,7 @@
             return ret;
         }
     } else
-#endif /* MBEDTLS_ECDH_C */
+#endif /* MBEDTLS_PK_CAN_ECDH */
     if (0 /* other KEMs? */) {
         /* Do something */
     } else {